Create app.py

app.py

@@ -0,0 +1,143 @@
+import gradio as gr
+import numpy as np
+import torch
+from diffusers import StableDiffusionInpaintPipeline
+from PIL import Image
+from segment_anything import SamPredictor, sam_model_registry, SamAutomaticMaskGenerator
+from diffusers import ControlNetModel
+from diffusers import UniPCMultistepScheduler
+from controlnet_inpaint import StableDiffusionControlNetInpaintPipeline
+import colorsys
+
+sam_checkpoint = "weights/sam_vit_h_4b8939.pth"
+model_type = "vit_h"
+device = "cuda"
+
+sam = sam_model_registry[model_type](checkpoint=sam_checkpoint)
+sam.to(device=device)
+predictor = SamPredictor(sam)
+mask_generator = SamAutomaticMaskGenerator(sam)
+
+# pipe = StableDiffusionInpaintPipeline.from_pretrained(
+#     "stabilityai/stable-diffusion-2-inpainting",
+#     torch_dtype=torch.float16,
+# )
+# pipe = pipe.to("cuda")
+
+controlnet = ControlNetModel.from_pretrained(
+    "lllyasviel/sd-controlnet-seg",
+    torch_dtype=torch.float16,
+)
+pipe = StableDiffusionControlNetInpaintPipeline.from_pretrained(
+    "runwayml/stable-diffusion-inpainting",
+    controlnet=controlnet,
+    torch_dtype=torch.float16,
+)
+pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
+pipe.enable_model_cpu_offload()
+pipe.enable_xformers_memory_efficient_attention()
+
+
+with gr.Blocks() as demo:
+    selected_pixels = gr.State([])
+    with gr.Row():
+        input_img = gr.Image(label="Input")
+        mask_img = gr.Image(label="Mask")
+        seg_img = gr.Image(label="Segmentation")
+        output_img = gr.Image(label="Output")
+
+    with gr.Row():
+        prompt_text = gr.Textbox(lines=1, label="Prompt")
+        negative_prompt_text = gr.Textbox(lines=1, label="Negative Prompt")
+        is_background = gr.Checkbox(label="Background")
+
+    with gr.Row():
+        submit = gr.Button("Submit")
+        clear = gr.Button("Clear")
+
+    def generate_mask(image, bg, sel_pix, evt: gr.SelectData):
+        sel_pix.append(evt.index)
+        predictor.set_image(image)
+        input_point = np.array(sel_pix)
+        input_label = np.ones(input_point.shape[0])
+        mask, _, _ = predictor.predict(
+            point_coords=input_point,
+            point_labels=input_label,
+            multimask_output=False,
+        )
+        if bg:
+            mask = np.logical_not(mask)
+        mask = Image.fromarray(mask[0, :, :])
+        segs = mask_generator.generate(image)
+        boolean_masks = [s["segmentation"] for s in segs]
+        finseg = np.zeros((boolean_masks[0].shape[0], boolean_masks[0].shape[1], 3), dtype=np.uint8)
+        # Loop over the boolean masks and assign a unique color to each class
+        for class_id, boolean_mask in enumerate(boolean_masks):
+            hue = class_id * 1.0 / len(boolean_masks)
+            rgb = tuple(int(i * 255) for i in colorsys.hsv_to_rgb(hue, 1, 1))
+            rgb_mask = np.zeros((boolean_mask.shape[0], boolean_mask.shape[1], 3), dtype=np.uint8)
+            rgb_mask[:, :, 0] = boolean_mask * rgb[0]
+            rgb_mask[:, :, 1] = boolean_mask * rgb[1]
+            rgb_mask[:, :, 2] = boolean_mask * rgb[2]
+            finseg += rgb_mask
+
+        return mask, finseg
+
+    def inpaint(image, mask, seg_img, prompt, negative_prompt):
+        image = Image.fromarray(image)
+        mask = Image.fromarray(mask)
+        seg_img = Image.fromarray(seg_img)
+
+        image = image.resize((512, 512))
+        mask = mask.resize((512, 512))
+        seg_img = seg_img.resize((512, 512))
+
+        output = pipe(prompt, image, mask, seg_img, negative_prompt=negative_prompt).images[0]
+        return output
+
+    def _clear(sel_pix, img, mask, seg, out, prompt, neg_prompt, bg):
+        sel_pix = []
+        img = None
+        mask = None
+        seg = None
+        out = None
+        prompt = ""
+        neg_prompt = ""
+        bg = False
+        return img, mask, seg, out, prompt, neg_prompt, bg
+
+    input_img.select(
+        generate_mask,
+        [input_img, is_background, selected_pixels],
+        [mask_img, seg_img],
+    )
+    submit.click(
+        inpaint,
+        inputs=[input_img, mask_img, seg_img, prompt_text, negative_prompt_text],
+        outputs=[output_img],
+    )
+    clear.click(
+        _clear,
+        inputs=[
+            selected_pixels,
+            input_img,
+            mask_img,
+            seg_img,
+            output_img,
+            prompt_text,
+            negative_prompt_text,
+            is_background,
+        ],
+        outputs=[
+            input_img,
+            mask_img,
+            seg_img,
+            output_img,
+            prompt_text,
+            negative_prompt_text,
+            is_background,
+        ],
+    )
+
+if __name__ == "__main__":
+    demo.queue(concurrency_count=50).launch()