Update app.py

app.py

@@ -1,124 +1,96 @@
-from __future__ import annotations
-import math
-import random
 import gradio as gr
+from gradio_imageslider import ImageSlider
+from PIL import Image, ImageDraw, ImageFont
 import numpy as np
+import cv2
 import torch
-from PIL import Image
-from diffusers import StableDiffusionXLImg2ImgPipeline, EDMEulerScheduler, AutoencoderKL
-from huggingface_hub import hf_hub_download
+from torchvision import transforms
+from transformers import AutoModelForImageSegmentation
 
-vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16)
+torch.set_float32_matmul_precision(["high", "highest"][0])
 
-pipe_edit = StableDiffusionXLImg2ImgPipeline.from_single_file(
-    hf_hub_download(repo_id="stabilityai/cosxl", filename="cosxl_edit.safetensors"),
-    num_in_channels=8,
-    is_cosxl_edit=True,
-    vae=vae,
-    torch_dtype=torch.float16,
+# Load BiRefNet model for background removal
+birefnet = AutoModelForImageSegmentation.from_pretrained(
+    "ZhengPeng7/BiRefNet", trust_remote_code=True
 )
-
-pipe_edit.scheduler = EDMEulerScheduler(sigma_min=0.002, sigma_max=120.0, sigma_data=1.0, prediction_type="v_prediction")
-pipe_edit.to("cuda")
-
-refiner = StableDiffusionXLImg2ImgPipeline.from_pretrained(
-    "stabilityai/stable-diffusion-xl-refiner-1.0", 
-    vae=vae, 
-    torch_dtype=torch.float16, 
-    use_safetensors=True, 
-    variant="fp16"
+birefnet.to("cuda")
+transform_image = transforms.Compose(
+    [
+        transforms.Resize((1024, 1024)),
+        transforms.ToTensor(),
+        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
+    ]
 )
-refiner.to("cuda")
 
-def set_timesteps_patched(self, num_inference_steps: int, device=None):
-    self.num_inference_steps = num_inference_steps
-    ramp = np.linspace(0, 1, self.num_inference_steps)
-    sigmas = torch.linspace(math.log(self.config.sigma_min), math.log(self.config.sigma_max), len(ramp)).exp().flip(0)
-    sigmas = sigmas.to(dtype=torch.float32, device=device)
-    self.timesteps = self.precondition_noise(sigmas)
-    self.sigmas = torch.cat([sigmas, torch.zeros(1, device=sigmas.device)])
-    self._step_index = None
-    self._begin_index = None
-    self.sigmas = self.sigmas.to("cpu")
+def load_img(image, output_type="numpy"):
+    if output_type == "pil":
+        return Image.open(image).convert("RGB")
+    else:
+        return np.array(Image.open(image).convert("RGB"))
+
+def add_text_to_image(image, text, position, color, font_size):
+    img = Image.fromarray(image)
+    draw = ImageDraw.Draw(img)
+    font = ImageFont.truetype("arial.ttf", font_size)
+    draw.text(position, text, fill=color, font=font)
+    return np.array(img)
 
-EDMEulerScheduler.set_timesteps = set_timesteps_patched
+def inpaint_image(image, mask, inpaint_radius):
+    img_cv = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
+    mask_cv = cv2.cvtColor(mask, cv2.COLOR_RGB2GRAY)
+    result = cv2.inpaint(img_cv, mask_cv, inpaint_radius, cv2.INPAINT_TELEA)
+    return cv2.cvtColor(result, cv2.COLOR_BGR2RGB)
 
-def king(input_image, instruction: str, negative_prompt: str = "", steps: int = 25, randomize_seed: bool = True, seed: int = 2404, guidance_scale: float = 6, progress=gr.Progress(track_tqdm=True)):
-    input_image = Image.open(input_image).convert('RGB')
-    if randomize_seed:
-        seed = random.randint(0, 999999)
-    generator = torch.manual_seed(seed)
-    output_image = pipe_edit(
-        instruction,
-        negative_prompt=negative_prompt,
-        image=input_image,
-        guidance_scale=guidance_scale,
-        image_guidance_scale=1.5,
-        width=input_image.width,
-        height=input_image.height,
-        num_inference_steps=steps,
-        generator=generator,
-        output_type="latent",
-    ).images
-    refine = refiner(
-        prompt=f"{instruction}, 4k, hd, high quality, masterpiece",
-        negative_prompt=negative_prompt,
-        guidance_scale=7.5,
-        num_inference_steps=steps,
-        image=output_image,
-        generator=generator,
-    ).images[0]  
-    return seed, refine
+def background_removal(image):
+    im = load_img(image, output_type="pil")
+    im = im.convert("RGB")
+    image_size = im.size
+    origin = im.copy()
+    image = load_img(im)
+    input_images = transform_image(image).unsqueeze(0).to("cuda")
+    with torch.no_grad():
+        preds = birefnet(input_images)[-1].sigmoid().cpu()
+    pred = preds[0].squeeze()
+    pred_pil = transforms.ToPILImage()(pred)
+    mask = pred_pil.resize(image_size)
+    im.putalpha(mask)
+    return (im, origin)
 
-css = '''
-.gradio-container{max-width: 700px !important}
-h1{text-align:center}
-footer {
-    visibility: hidden
-}
-'''
+def update_image(image, text, color, font_size, mask_image, inpaint_radius):
+    img_with_text = add_text_to_image(image, text, (50, 50), color, font_size)
+    if mask_image is not None:
+        mask = np.array(mask_image)
+        img_with_text = inpaint_image(img_with_text, mask, inpaint_radius)
+    return img_with_text
 
-examples = [
-    ["./supercar.png", "make it red"],
-    ["./red_car.png", "add some snow"],
-]
+def fn(image):
+    return background_removal(image)
 
-with gr.Blocks(css=css) as demo:
-    gr.Markdown("# Image Editing\n### Note: First image generation takes time")
-    with gr.Row():
-        instruction = gr.Textbox(lines=1, label="Instruction", interactive=True)
-        generate_button = gr.Button("Run", scale=0)
-        
-    with gr.Row():
-        input_image = gr.Image(label="Image", type='filepath', interactive=True)
+slider1 = ImageSlider(label="Original Image", type="pil")
+slider2 = ImageSlider(label="Processed Image", type="pil")
 
-    with gr.Row():
-        guidance_scale = gr.Number(value=6.0, step=0.1, label="Guidance Scale", interactive=True)
-        steps = gr.Number(value=25, step=1, label="Steps", interactive=True)
+image_input = gr.Image(label="Upload an image for background removal")
+text_input = gr.Textbox(label="Enter Text to Add", placeholder="Your text here...")
+color_input = gr.ColorPicker(label="Text Color")
+font_size_input = gr.Slider(minimum=10, maximum=100, label="Font Size")
+mask_input = gr.Image(type="numpy", label="Upload Mask Image (for Inpainting)", optional=True)
+inpaint_radius_input = gr.Slider(minimum=1, maximum=50, value=3, label="Inpaint Radius")
 
-    with gr.Accordion("Advanced options", open=False):
-        with gr.Row():
-            negative_prompt = gr.Text(
-                label="Negative prompt",
-                max_lines=1,
-                value="(deformed, distorted, disfigured:1.3), poorly drawn, bad anatomy, wrong anatomy, extra limb, missing limb, floating limbs, (mutated hands and fingers:1.4), disconnected limbs, ugly, disgusting, blurry, amputation,(face asymmetry, eyes asymmetry, deformed eyes, open mouth)",
-                visible=True
-            )
-        with gr.Row():
-            randomize_seed = gr.Checkbox(label="Randomize Seed", value=True, interactive=True)
-            seed = gr.Number(value=2404, step=1, label="Seed", interactive=True)
+bg_removal_interface = gr.Interface(
+    fn, inputs=image_input, outputs=slider1, examples=["chameleon.jpg"]
+)
 
-    gr.Examples(
-        examples=examples,
-        inputs=[input_image, instruction],
-        outputs=[input_image],
-        cache_examples=False,
-    )
+design_editing_interface = gr.Interface(
+    fn=lambda image, text, color, font_size, mask_image, inpaint_radius: update_image(image, text, color, font_size, mask_image, inpaint_radius),
+    inputs=[image_input, text_input, color_input, font_size_input, mask_input, inpaint_radius_input],
+    outputs=slider2
+)
 
-    generate_button.click(
-        king,
-        inputs=[input_image, instruction, negative_prompt, steps, randomize_seed, seed, guidance_scale],
-        outputs=[seed, input_image],
-    )
+demo = gr.TabbedInterface(
+    [bg_removal_interface, design_editing_interface],
+    ["Background Removal", "Design Editing"],
+    title="Advanced Image Editor"
+)
 
-demo.queue(max_size=500).launch()
+if __name__ == "__main__":
+    demo.launch()