Update app.py

app.py

@@ -3,75 +3,56 @@ import torch
 from diffusers import StableDiffusionImg2ImgPipeline
 from torchvision import transforms
 from PIL import Image
-import io
-import base64
-import spaces
+import io, base64, spaces
 
-# Load Ghibli model
+# Ghibli model
 pipe_ghibli = StableDiffusionImg2ImgPipeline.from_pretrained(
     "nitrosocke/Ghibli-Diffusion", torch_dtype=torch.float16
 ).to("cuda")
 
-# Load CartoonGAN model
-cartoon_model = torch.hub.load('AK391/CartoonGAN', 'cartoon_gan', pretrained=True).to("cuda").eval()
+# CartoonGAN model via torch.hub
+cartoon_model = torch.hub.load(
+    'znxlwm/pytorch-CartoonGAN', 'CartoonGAN',
+    pretrained=True, trust_repo=True
+).to("cuda").eval()
 
-# Base64 utilities
-def pil_to_base64(img: Image.Image) -> str:
-    buffer = io.BytesIO()
-    img.save(buffer, format="PNG")
-    return base64.b64encode(buffer.getvalue()).decode()
-
-def base64_to_pil(b64: str) -> Image.Image:
-    image_data = base64.b64decode(b64)
-    return Image.open(io.BytesIO(image_data)).convert("RGB")
+# 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_gan(img: Image.Image) -> Image.Image:
-    transform = transforms.Compose([
-        transforms.Resize((256, 256)),
-        transforms.ToTensor()
-    ])
-    img_tensor = transform(img).unsqueeze(0).to("cuda")
-    with torch.no_grad():
-        output = cartoon_model(img_tensor)[0].clamp(0, 1).cpu()
-    output_pil = transforms.ToPILImage()(output)
-    return output_pil
-
-# Unified effect processor
-def process_image(input_image: Image.Image, effect: str) -> Image.Image:
-    if effect == "ghibli":
-        output_image = pipe_ghibli(prompt="ghibli style", image=input_image, strength=0.5, guidance_scale=7.5).images[0]
-    else:
-        output_image = apply_cartoon_gan(input_image)
-    return output_image
+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(input_b64: str, effect: str) -> str:
-    input_image = base64_to_pil(input_b64)
-    output_image = process_image(input_image, effect)
-    return pil_to_base64(output_image)
+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")
-
+    gr.Markdown("# 🎨 Ghibli & CartoonGAN Effects (ZeroGPU)")
     with gr.Tab("Web UI"):
-        with gr.Row():
-            image_input = gr.Image(type="pil", label="Upload Image")
-            effect_selector = gr.Radio(["ghibli", "cartoon"], label="Select Effect")
-        with gr.Row():
-            apply_button = gr.Button("Apply Effect")
-        with gr.Row():
-            image_output = gr.Image(label="Processed Image")
-
-        apply_button.click(process_image, [image_input, effect_selector], image_output)
-
+        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"):
-        base64_input = gr.Textbox(label="Input Image (Base64)", lines=5)
-        effect_choice = gr.Radio(["ghibli", "cartoon"], label="Select Effect")
-        api_button = gr.Button("Run API")
-        base64_output = gr.Textbox(label="Output Image (Base64)", lines=5)
-
-        api_button.click(process_base64, [base64_input, effect_choice], base64_output)
+        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()