Update app.py

app.py

@@ -1,62 +1,77 @@
-# app.py
 import gradio as gr
-import torch, io, base64
-from PIL import Image
+import torch
 from diffusers import StableDiffusionImg2ImgPipeline
-from vtoonify_model import load_vtoonify  # see below
+from torchvision import transforms
+from PIL import Image
+import io
+import base64
 import spaces
 
-# Load models
+# Load Ghibli model
 pipe_ghibli = StableDiffusionImg2ImgPipeline.from_pretrained(
     "nitrosocke/Ghibli-Diffusion", torch_dtype=torch.float16
-).to("cuda")  # Ghibli-style fine-tuned SD :contentReference[oaicite:1]{index=1}
+).to("cuda")
 
-pipe_vtoonify = load_vtoonify().to("cuda")  # cartoonization model loader
+# Load CartoonGAN model
+cartoon_model = torch.hub.load('AK391/CartoonGAN', 'cartoon_gan', pretrained=True).to("cuda").eval()
 
-# Helpers for base64 conversion
-def pil_to_b64(img: Image.Image) -> str:
-    buf = io.BytesIO()
-    img.save(buf, format="PNG")
-    return base64.b64encode(buf.getvalue()).decode()
+# 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 b64_to_pil(b64: str) -> Image.Image:
-    data = base64.b64decode(b64)
-    return Image.open(io.BytesIO(data)).convert("RGB")
+def base64_to_pil(b64: str) -> Image.Image:
+    image_data = base64.b64decode(b64)
+    return Image.open(io.BytesIO(image_data)).convert("RGB")
 
-# Core processor
-def run_effect(input_b64: str, effect: str) -> dict:
-    img = b64_to_pil(input_b64)
+# 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":
-        out = pipe_ghibli(prompt="ghibli style", image=img, strength=0.5, guidance_scale=7.5).images[0]
+        output_image = pipe_ghibli(prompt="ghibli style", image=input_image, strength=0.5, guidance_scale=7.5).images[0]
     else:
-        out = pipe_vtoonify(img)
-    return {"output_b64": pil_to_b64(out)}
-
-@gr.utils.decorators.thread_safe()
-@spaces.GPU  # enables GPU on ZeroGPU Infra
-def api_process(input_b64, effect):
-    return run_effect(input_b64, effect)
+        output_image = apply_cartoon_gan(input_image)
+    return output_image
 
-def gradio_process(img: Image.Image, effect: str) -> Image.Image:
-    # Reuse logic, bypass base64
-    in_b64 = pil_to_b64(img)
-    return b64_to_pil(run_effect(in_b64, effect)["output_b64"])
+@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)
 
+# Gradio UI
 with gr.Blocks() as demo:
-    gr.Markdown("# Ghibli & VToonify Effects 🎨")
+    gr.Markdown("# 🎨 Ghibli & CartoonGAN Effects")
 
     with gr.Tab("Web UI"):
-        inp = gr.Image(type="pil", label="Upload Image")
-        eff = gr.Radio(["ghibli", "vtoonify"], label="Effect")
-        btn = gr.Button("Apply Effect")
-        out = gr.Image(label="Result")
-        btn.click(gradio_process, [inp, eff], out)
-
-    with gr.Tab("API (base64)"):
-        inp_b64 = gr.Textbox(lines=4, label="Input Image (base64)")
-        eff2 = gr.Radio(["ghibli", "vtoonify"], label="Effect")
-        btn2 = gr.Button("Run API")
-        out_b64 = gr.Textbox(lines=4, label="Output Image (base64)")
-        btn2.click(api_process, [inp_b64, eff2], out_b64)
+        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)
+
+    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)
 
 demo.launch()