Update app.py

app.py

@@ -1,85 +1,102 @@
 import gradio as gr
 import torch
-import numpy as np
-from diffusers import StableDiffusionXLInpaintPipeline
-from PIL import Image, ImageDraw
-from transformers import DetrImageProcessor, DetrForObjectDetection
+import base64
+import io
+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
 
-device = "cuda" if torch.cuda.is_available() else "cpu"
-
-# Load the Stable Diffusion XL Inpainting model
-pipe = StableDiffusionXLInpaintPipeline.from_pretrained(
-    "stabilityai/stable-diffusion-xl-base-1.0",
-    torch_dtype=torch.float16 if device == "cuda" else torch.float32,
-    variant="fp16",
-).to(device)
-
-# Load the DETR object detection model
-processor = DetrImageProcessor.from_pretrained("facebook/detr-resnet-50")
-detector = DetrForObjectDetection.from_pretrained("facebook/detr-resnet-50").to(device)
-
-@spaces.GPU
-def detect_and_replace_humans(input_image, prompt):
-    if input_image is None or prompt == "":
-        return None
-
-    image_np = np.array(input_image)
-    inputs = processor(images=input_image, return_tensors="pt").to(device)
-
-    outputs = detector(**inputs)
-    target_sizes = torch.tensor([image_np.shape[:2]]).to(device)
-
-    results = processor.post_process_object_detection(outputs, target_sizes=target_sizes, threshold=0.9)[0]
-
-    mask = Image.new("L", input_image.size, 0)
-    draw = ImageDraw.Draw(mask)
-
-    found = False
-    for score, label, box in zip(results["scores"], results["labels"], results["boxes"]):
-        if detector.config.id2label[label.item()] == "person":
-            box = [int(i) for i in box.tolist()]
-            draw.rectangle(box, fill=255)
-            found = True
-
-    if not found:
-        return "No human detected."
-
-    # Pre-defined positive and negative prompts
-    positive_prompt = (
-        "Replace the masked humans with imaginary Indian bride and groom wearing traditional Indian wedding attire, "
-        "with detailed embroidery, colorful saree and sherwani, realistic faces, natural skin texture, matching pose, "
-        "perfect lighting, and the same camera perspective. Keep the background unchanged."
-    )
+# 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]
 
-    negative_prompt = (
-        "blurry, distorted, deformed, double face, extra limbs, low quality, bad proportions, low resolution, "
-        "changed background, multiple faces, duplicate body parts, cartoon, watermark, text"
-    )
+    clear_cache(pipe.transformer)
 
-    # Inpainting process
-    output = pipe(
-        prompt=positive_prompt,
-        negative_prompt=negative_prompt,
-        image=input_image,
-        mask_image=mask,
-        num_inference_steps=40,
-        guidance_scale=8.5
-    ).images[0]
+    return image_to_base64(result)
 
-    return output
+# 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 UI
+# Gradio App
 with gr.Blocks() as demo:
-    gr.Markdown("## Replace Humans with Imaginary Indian Bride and Groom (Background Preserved)")
+    gr.Markdown("# 🎨 Ghibli Style Cartoonizer using EasyControl")
 
     with gr.Row():
-        input_image = gr.Image(type="pil", label="Input Image")
-        output_image = gr.Image(type="pil", label="Output Image")
-
-    prompt_text = gr.Textbox(label="Prompt (Optional, Predefined Prompt Used)", placeholder="You can leave this blank")
-    submit = gr.Button("Submit")
+        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
+    )
 
-    submit.click(detect_and_replace_humans, inputs=[input_image, prompt_text], 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"
+    )
 
 demo.launch()