Update app.py

app.py

@@ -1,102 +1,61 @@
+# app.py
 import gradio as gr
-import torch
-import base64
-import io
+import torch, io, base64
 from PIL import Image
-from diffusers import StableDiffusionPipeline
-from safetensors.torch import load_file
-from src.pipeline import FluxPipeline
-from src.transformer_flux import FluxTransformer2DModel
-from src.lora_helper import set_single_lora, clear_cache
-import spaces
+from diffusers import StableDiffusionImg2ImgPipeline
+from vtoonify_model import load_vtoonify  # see below
+
+# Load models
+pipe_ghibli = StableDiffusionImg2ImgPipeline.from_pretrained(
+    "nitrosocke/Ghibli-Diffusion", torch_dtype=torch.float16
+).to("cuda")  # Ghibli-style fine-tuned SD :contentReference[oaicite:1]{index=1}
+
+pipe_vtoonify = load_vtoonify().to("cuda")  # cartoonization model loader
+
+# 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()
+
+def b64_to_pil(b64: str) -> Image.Image:
+    data = base64.b64decode(b64)
+    return Image.open(io.BytesIO(data)).convert("RGB")
+
+# Core processor
+def run_effect(input_b64: str, effect: str) -> dict:
+    img = b64_to_pil(input_b64)
+    if effect == "ghibli":
+        out = pipe_ghibli(prompt="ghibli style", image=img, 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)
+
+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"])
 
-# Load Base Model and LoRA
-base_model = "black-forest-labs/FLUX.1-dev"
-lora_path = "checkpoints/models/Ghibli.safetensors"
-
-# Load the main pipeline
-pipe = FluxPipeline.from_pretrained(base_model, torch_dtype=torch.float16)
-transformer = FluxTransformer2DModel.from_pretrained(base_model, subfolder="transformer", torch_dtype=torch.float16)
-pipe.transformer = transformer
-pipe.to("cuda")
-
-# Load LoRA
-set_single_lora(pipe.transformer, lora_path, lora_weights=[1], cond_size=512)
-
-# Base64 to Image
-def base64_to_image(base64_str):
-    image_data = base64.b64decode(base64_str)
-    return Image.open(io.BytesIO(image_data)).convert("RGB")
-
-# Image to Base64
-def image_to_base64(image):
-    buffered = io.BytesIO()
-    image.save(buffered, format="PNG")
-    return base64.b64encode(buffered.getvalue()).decode()
-
-# Cartoonizer function
-def cartoonize_base64(b64_image, prompt="Ghibli Studio style, hand-drawn anime illustration", height=768, width=768, seed=42):
-    input_image = base64_to_image(b64_image)
-
-    generator = torch.Generator(device="cuda").manual_seed(int(seed))
-
-    result = pipe(
-        prompt=prompt,
-        height=int(height),
-        width=int(width),
-        guidance_scale=3.5,
-        num_inference_steps=25,
-        generator=generator,
-        spatial_images=[input_image],
-        cond_size=512
-    ).images[0]
-
-    clear_cache(pipe.transformer)
-
-    return image_to_base64(result)
-
-# Gradio UI function
-def ui_cartoonize(image, prompt, height, width, seed):
-    buffered = io.BytesIO()
-    image.save(buffered, format="PNG")
-    b64_image = base64.b64encode(buffered.getvalue()).decode()
-    cartoon_b64 = cartoonize_base64(b64_image, prompt, height, width, seed)
-    cartoon_image = base64_to_image(cartoon_b64)
-    return cartoon_image
-
-# Gradio App
 with gr.Blocks() as demo:
-    gr.Markdown("# 🎨 Ghibli Style Cartoonizer using EasyControl")
-
-    with gr.Row():
-        with gr.Column():
-            input_image = gr.Image(type="pil", label="Upload Image")
-            prompt = gr.Textbox(label="Prompt", value="Ghibli Studio style, hand-drawn anime illustration")
-            height = gr.Slider(512, 1024, step=64, value=768, label="Height")
-            width = gr.Slider(512, 1024, step=64, value=768, label="Width")
-            seed = gr.Number(label="Seed", value=42)
-            generate_btn = gr.Button("Generate Ghibli Image")
-        with gr.Column():
-            output_image = gr.Image(label="Cartoonized Output")
-
-    generate_btn.click(
-        fn=ui_cartoonize,
-        inputs=[input_image, prompt, height, width, seed],
-        outputs=output_image
-    )
-
-    # Gradio API: Accept base64, return base64
-    gr.Interface(
-        fn=cartoonize_base64,
-        inputs=[
-            gr.Text(label="Base64 Image Input"),
-            gr.Text(label="Prompt"),
-            gr.Number(label="Height", value=768),
-            gr.Number(label="Width", value=768),
-            gr.Number(label="Seed", value=42)
-        ],
-        outputs=gr.Text(label="Base64 Cartoon Output"),
-        api_name="predict"
-    )
+    gr.Markdown("# Ghibli & VToonify 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)
 
 demo.launch()