Update app.py

app.py

@@ -66,59 +66,27 @@ def resize_image(input_image, resolution):
     img = cv2.resize(input_image, (W, H), interpolation=cv2.INTER_LANCZOS4 if k > 1 else cv2.INTER_AREA)
     return img
 
-
-def process(input_image, 
-    original_image, 
-    original_mask, 
-    input_mask, 
-    selected_points, 
-    prompt, 
-    negative_prompt, 
-    blended, 
-    invert_mask, 
-    control_strength, 
-    seed, 
-    randomize_seed, 
-    guidance_scale, 
-    num_inference_steps):
-    
+def process(original_image, input_mask, prompt, negative_prompt, blended, invert_mask, control_strength, seed, randomize_seed, guidance_scale, num_inference_steps):
     if original_image is None:
-        if input_image:
-            original_image = input_image
-            original_mask = input_mask
-        else:
-            raise gr.Error('Please upload the input image')
-    if (original_mask is None or len(selected_points)==0) and input_mask is None:
-        raise gr.Error("Please click the region where you hope unchanged/changed, or upload a white-black Mask image")
+        raise gr.Error('Please upload the input image')
+    if input_mask is None:
+        raise gr.Error("Please upload a white-black Mask image")
     
-    # load example image
-    if isinstance(original_image, int):
-        image_name = image_examples[original_image][0]
-        original_image = cv2.imread(image_name)
-        original_image = cv2.cvtColor(original_image, cv2.COLOR_BGR2RGB)
-
-    if input_mask is not None:
-        H,W=original_image.shape[:2]
-        original_mask = cv2.resize(input_mask, (W, H))
-    else:
-        original_mask = np.clip(255 - original_mask, 0, 255).astype(np.uint8)
+    H, W = original_image.shape[:2]
+    original_mask = cv2.resize(input_mask, (W, H))
 
     if invert_mask:
-        original_mask=255-original_mask
-
-    mask = 1.*(original_mask.sum(-1)>255)[:,:,np.newaxis]
-    masked_image = original_image * (1-mask)
-
+        original_mask = 255 - original_mask
+    mask = 1.*(original_mask.sum(-1) > 255)[:,:,np.newaxis]
+    masked_image = original_image * (1 - mask)
     init_image = Image.fromarray(masked_image.astype(np.uint8)).convert("RGB")
     mask_image = Image.fromarray(original_mask.astype(np.uint8)).convert("RGB")
-
-    generator = torch.Generator("cuda").manual_seed(random.randint(0,2147483647) if randomize_seed else seed)
-
+    generator = torch.Generator("cuda").manual_seed(random.randint(0, 2147483647) if randomize_seed else seed)
     image = pipe(
-        [prompt]*2, 
-        init_image, 
-        mask_image, 
-        num_inference_steps=num_inference_steps, 
+        [prompt]*2,
+        init_image,
+        mask_image,
+        num_inference_steps=num_inference_steps,
         guidance_scale=guidance_scale,
         generator=generator,
         brushnet_conditioning_scale=float(control_strength),
@@ -126,220 +94,43 @@ def process(input_image,
     ).images
 
     if blended:
-        if control_strength<1.0:
+        if control_strength < 1.0:
             raise gr.Error('Using blurred blending with control strength less than 1.0 is not allowed')
-        blended_image=[]
-        # blur, you can adjust the parameters for better performance
+        blended_image = []
         mask_blurred = cv2.GaussianBlur(mask*255, (21, 21), 0)/255
         mask_blurred = mask_blurred[:,:,np.newaxis]
-        mask = 1-(1-mask) * (1-mask_blurred)
+        mask = 1 - (1 - mask) * (1 - mask_blurred)
         for image_i in image:
-            image_np=np.array(image_i)
-            image_pasted=original_image * (1-mask) + image_np*mask
-
-            image_pasted=image_pasted.astype(image_np.dtype)
+            image_np = np.array(image_i)
+            image_pasted = original_image * (1 - mask) + image_np * mask
+            image_pasted = image_pasted.astype(image_np.dtype)
             blended_image.append(Image.fromarray(image_pasted))
-        
-        image=blended_image
+        image = blended_image
 
     return image
 
-block = gr.Blocks(
-        theme=gr.themes.Soft(
-             radius_size=gr.themes.sizes.radius_none,
-             text_size=gr.themes.sizes.text_md
-         )
-        ).queue()
-with block:
+# Create Gradio interface
+with gr.Blocks() as demo:
     with gr.Row():
         with gr.Column():
-            
-            gr.HTML(f"""
-                    <div style="text-align: center;">
-                        <h1>BrushNet: A Plug-and-Play Image Inpainting Model with Decomposed Dual-Branch Diffusion</h1>
-                        <div style="display: flex; justify-content: center; align-items: center; text-align: center;">
-                            <a href=""></a>
-                            <a href='https://tencentarc.github.io/BrushNet/'><img src='https://img.shields.io/badge/Project_Page-BrushNet-green' alt='Project Page'></a>
-                            <a href='https://arxiv.org/abs/2403.06976'><img src='https://img.shields.io/badge/Paper-Arxiv-blue'></a>
-                        </div>
-                        </br>
-                    </div>
-            """)
-
-
-    with gr.Accordion(label="🧭 Instructions:", open=True, elem_id="accordion"):
-        with gr.Row(equal_height=True):
-            gr.Markdown("""
-            - ⭐️ <b>step1: </b>Upload or select one image from Example
-            - ⭐️ <b>step2: </b>Click on Input-image to select the object to be retained (or upload a white-black Mask image, in which white color indicates the region you want to keep unchanged). You can tick the 'Invert Mask' box to switch region unchanged and change.
-            - ⭐️ <b>step3: </b>Input prompt for generating new contents
-            - ⭐️ <b>step4: </b>Click Run button
-            """)                          
-    with gr.Row():
-        with gr.Column():
-            with gr.Column(elem_id="Input"):
-                with gr.Row():
-                    with gr.Tabs(elem_classes=["feedback"]):
-                        with gr.TabItem("Input Image"):
-                            input_image = gr.Image(type="numpy", label="input",scale=2, height=1024)
-                original_image = gr.State(value=None, label="index")
-                original_mask = gr.State(value=None)
-                selected_points = gr.State([],label="select points")
-                with gr.Row(elem_id="Seg"):
-                    radio = gr.Radio(['foreground', 'background'], label='Click to seg: ', value='foreground',scale=2)
-                    undo_button = gr.Button('Undo seg', elem_id="btnSEG",scale=1)
-            prompt = gr.Textbox(label="Prompt", placeholder="Please input your prompt",value='',lines=1)
-            negative_prompt = gr.Text(
-                        label="Negative Prompt",
-                        max_lines=5,
-                        placeholder="Please input your negative prompt",
-                        value='ugly, low quality',lines=1
-                    )
-            with gr.Group():
-                with gr.Row():
-                    blending = gr.Checkbox(label="Blurred Blending", value=False)
-                    invert_mask = gr.Checkbox(label="Invert Mask", value=True)
-            run_button = gr.Button("Run",elem_id="btn")
-            
-            with gr.Accordion("More input params (highly-recommended)", open=False, elem_id="accordion1"):
-                control_strength = gr.Slider(
-                    label="Control Strength: ", show_label=True, minimum=0, maximum=1.1, value=1, step=0.01
-                    )
-                with gr.Group():
-                    seed = gr.Slider(
-                        label="Seed: ", minimum=0, maximum=2147483647, step=1, value=551793204
-                    )
-                    randomize_seed = gr.Checkbox(label="Randomize seed", value=False)
-                
-                with gr.Group():
-                    with gr.Row():
-                        guidance_scale = gr.Slider(
-                            label="Guidance scale",
-                            minimum=1,
-                            maximum=12,
-                            step=0.1,
-                            value=12,
-                        )
-                        num_inference_steps = gr.Slider(
-                            label="Number of inference steps",
-                            minimum=1,
-                            maximum=50,
-                            step=1,
-                            value=50,
-                        )
-                with gr.Row(elem_id="Image"):
-                    with gr.Tabs(elem_classes=["feedback1"]):
-                        with gr.TabItem("User-specified Mask Image (Optional)"):
-                            input_mask = gr.Image(type="numpy", label="Mask Image", height=1024)
-            
+            original_image = gr.Image(type="numpy", label="Original Image")
+            input_mask = gr.Image(type="numpy", label="Mask Image")
+            prompt = gr.Textbox(label="Prompt")
+            negative_prompt = gr.Textbox(label="Negative Prompt", value='ugly, low quality')
+            blended = gr.Checkbox(label="Blurred Blending", value=False)
+            invert_mask = gr.Checkbox(label="Invert Mask", value=False)
+            control_strength = gr.Slider(label="Control Strength", minimum=0, maximum=1.1, value=1, step=0.01)
+            seed = gr.Slider(label="Seed", minimum=0, maximum=2147483647, step=1, value=551793204)
+            randomize_seed = gr.Checkbox(label="Randomize Seed", value=False)
+            guidance_scale = gr.Slider(label="Guidance Scale", minimum=1, maximum=12, step=0.1, value=7.5)
+            num_inference_steps = gr.Slider(label="Number of Inference Steps", minimum=1, maximum=50, step=1, value=50)
+            run_button = gr.Button("Run")
+        
         with gr.Column():
-            with gr.Tabs(elem_classes=["feedback"]):
-                with gr.TabItem("Outputs"):
-                    result_gallery = gr.Gallery(label='Output', show_label=False, elem_id="gallery", preview=True)
-    with gr.Row():
-        def process_example(input_image, prompt, input_mask, original_image, selected_points,result_gallery): #
-            return input_image, prompt, input_mask, original_image, [], result_gallery
-        example = gr.Examples(
-            label="Input Example",
-            examples=image_examples,
-            inputs=[input_image, prompt, input_mask, original_image, selected_points, result_gallery],
-            outputs=[input_image, prompt, input_mask, original_image, selected_points],
-            fn=process_example,
-            run_on_click=True,
-            examples_per_page=10
-        )
-
-    # once user upload an image, the original image is stored in `original_image`
-    def store_img(img):
-        # image upload is too slow
-        if min(img.shape[0], img.shape[1]) > 1024:
-            img = resize_image(img, 1024)
-        if max(img.shape[0], img.shape[1])*1.0/min(img.shape[0], img.shape[1])>2.0:
-            raise gr.Error('image aspect ratio cannot be larger than 2.0')
-        return img, img, [], None  # when new image is uploaded, `selected_points` should be empty
-
-    input_image.upload(
-        store_img,
-        [input_image],
-        [input_image, original_image, selected_points]
-    )
-
-    # user click the image to get points, and show the points on the image
-    def segmentation(img, sel_pix):
-        # online show seg mask
-        points = []
-        labels = []
-        for p, l in sel_pix:
-            points.append(p)
-            labels.append(l)
-        mobile_predictor.set_image(img if isinstance(img, np.ndarray) else np.array(img))
-        with torch.no_grad():
-            masks, _, _ = mobile_predictor.predict(point_coords=np.array(points), point_labels=np.array(labels), multimask_output=False)
-
-        output_mask = np.ones((masks.shape[1], masks.shape[2], 3))*255
-        for i in range(3):
-                output_mask[masks[0] == True, i] = 0.0
-
-        mask_all = np.ones((masks.shape[1], masks.shape[2], 3))
-        color_mask = np.random.random((1, 3)).tolist()[0]
-        for i in range(3):
-            mask_all[masks[0] == True, i] = color_mask[i]
-        masked_img = img / 255 * 0.3 + mask_all * 0.7
-        masked_img = masked_img*255
-        ## draw points
-        for point, label in sel_pix:
-            cv2.drawMarker(masked_img, point, colors[label], markerType=markers[label], markerSize=20, thickness=5)
-        return masked_img, output_mask
-    
-    def get_point(img, sel_pix, point_type, evt: gr.SelectData):
-        if point_type == 'foreground':
-            sel_pix.append((evt.index, 1))   # append the foreground_point
-        elif point_type == 'background':
-            sel_pix.append((evt.index, 0))    # append the background_point
-        else:
-            sel_pix.append((evt.index, 1))    # default foreground_point
-
-        if isinstance(img, int):
-            image_name = image_examples[img][0]
-            img = cv2.imread(image_name)
-            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
-
-        # online show seg mask
-        masked_img, output_mask = segmentation(img, sel_pix)
-        return masked_img.astype(np.uint8), output_mask
-    
-    input_image.select(
-        get_point,
-        [original_image, selected_points, radio],
-        [input_image, original_mask],
-    )
-
-    # undo the selected point
-    def undo_points(orig_img, sel_pix):
-        # draw points
-        output_mask = None
-        if len(sel_pix) != 0:
-            if isinstance(orig_img, int):   # if orig_img is int, the image if select from examples
-                temp = cv2.imread(image_examples[orig_img][0])
-                temp = cv2.cvtColor(temp, cv2.COLOR_BGR2RGB)
-            else:
-                temp = orig_img.copy()
-            sel_pix.pop()
-            # online show seg mask
-            if len(sel_pix) !=0:
-                temp, output_mask = segmentation(temp, sel_pix)
-            return temp.astype(np.uint8), output_mask
-        else:
-            gr.Error("Nothing to Undo")
-    
-    undo_button.click(
-        undo_points,
-        [original_image, selected_points],
-        [input_image, original_mask]
-    )
-
-    ips=[input_image, original_image, original_mask, input_mask, selected_points, prompt, negative_prompt, blending, invert_mask, control_strength, seed, randomize_seed, guidance_scale, num_inference_steps]
-    run_button.click(fn=process, inputs=ips, outputs=[result_gallery])
+            result_gallery = gr.Gallery(label='Output', show_label=False, elem_id="gallery", preview=True)
 
+    inputs = [original_image, input_mask, prompt, negative_prompt, blended, invert_mask, control_strength, seed, randomize_seed, guidance_scale, num_inference_steps]
+    run_button.click(fn=process, inputs=inputs, outputs=[result_gallery])
 
-block.launch()
+demo.queue(concurrency_count=1, max_size=1, api_open=True)
+demo.launch(show_api=True)