Update app.py

app.py

@@ -1,58 +1,69 @@
 import gradio as gr
 import torch
-from diffusers import StableDiffusionImg2ImgPipeline
-from torchvision import transforms
+import io, base64
 from PIL import Image
-import io, base64, spaces
+from diffusers import StableDiffusionImg2ImgPipeline
+import spaces
 
-# Ghibli model
+# Load Ghibli model
 pipe_ghibli = StableDiffusionImg2ImgPipeline.from_pretrained(
-    "nitrosocke/Ghibli-Diffusion", torch_dtype=torch.float16
+    "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
+# Load AnimeGANv2 model
+animegan = torch.hub.load(
+    "bryandlee/animegan2-pytorch:main",
+    "generator",
+    pretrained="face_paint_512_v2"
 ).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")
+def pil_to_b64(img: Image.Image) -> str:
+    buf = io.BytesIO()
+    img.save(buf, format="PNG")
+    return base64.b64encode(buf.getvalue()).decode()
+
+def b64_to_pil(b64: str) -> Image.Image:
+    return Image.open(io.BytesIO(base64.b64decode(b64))).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)
+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")
+    with torch.no_grad():
+        out = animegan(tensor)[0].clamp(0,1).cpu()
+    return Image.fromarray((out.permute(1,2,0).numpy()*255).astype("uint8"))
 
-# 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)
+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]
+    else:
+        return apply_animegan(img)
 
 @spaces.GPU
-def process_base64(b64, effect):
+def process_base64(b64: str, effect: str) -> str:
     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)")
+    gr.Markdown("# 🎨 Ghibli & AnimeGANv2 (Cartoon) Effects")
+
     with gr.Tab("Web UI"):
         inp = gr.Image(type="pil")
-        eff = gr.Radio(["ghibli","cartoon"], label="Effect")
-        btn = gr.Button("Apply")
+        eff = gr.Radio(["ghibli", "anime"], label="Effect")
+        btn = gr.Button("Apply Effect")
         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")
+        b64_in = gr.Textbox(lines=5)
+        eff2 = gr.Radio(["ghibli", "anime"], label="Effect")
         btn2 = gr.Button("Run API")
-        out_b64 = gr.Textbox(lines=5)
-        btn2.click(process_base64, [in_b64, eff2], out_b64)
+        b64_out = gr.Textbox(lines=5)
+        btn2.click(process_base64, [b64_in, eff2], b64_out)
 
 demo.launch()