Update app.py

app.py

@@ -162,10 +162,45 @@ def process_dis(
         [path_out_fp32, path_out_vis],
     )
 
-def run_demo_server(pipe_depth, pipe_normal, pipe_dis):
+def process_matting(
+    pipe,
+    path_input,
+    processing_res=default_image_processing_res,
+):
+    name_base, name_ext = os.path.splitext(os.path.basename(path_input))
+    print(f"Processing image {name_base}{name_ext}")
+
+    path_output_dir = tempfile.mkdtemp()
+    path_out_fp32 = os.path.join(path_output_dir, f"{name_base}_matting_fp32.npy")
+    path_out_vis = os.path.join(path_output_dir, f"{name_base}_matting_colored.png")
+
+    input_image = Image.open(path_input)
+
+    pipe_out = pipe(
+        input_image,
+        processing_res=processing_res,
+        batch_size=1 if processing_res == 0 else 0,
+        show_progress_bar=False,
+        mode='seg',
+    )
+
+    depth_pred = pipe_out.pred_np
+    depth_colored = pipe_out.pred_colored
+
+    np.save(path_out_fp32, depth_pred)
+    depth_colored.save(path_out_vis)
+
+    return (
+        [path_out_vis],
+        [path_out_fp32, path_out_vis],
+    )
+
+    
+def run_demo_server(pipe_depth, pipe_normal, pipe_dis, pipe_matting):
     process_pipe_depth = spaces.GPU(functools.partial(process_depth, pipe_depth))
     process_pipe_normal = spaces.GPU(functools.partial(process_normal, pipe_normal))
     process_pipe_dis = spaces.GPU(functools.partial(process_dis, pipe_dis))
+    process_pipe_matting = spaces.GPU(functools.partial(process_matting, pipe_matting))
     gradio_theme = gr.themes.Default()
 
     with gr.Blocks(
@@ -409,6 +444,63 @@ def run_demo_server(pipe_depth, pipe_normal, pipe_dis):
                     cache_examples=False,
                 )
 
+            with gr.Tab("Matting"):
+                with gr.Row():
+                    with gr.Column():
+                        dis_image_input = gr.Image(
+                            label="Input Image",
+                            type="filepath",
+                        )
+                        with gr.Row():
+                            matting_image_submit_btn = gr.Button(
+                                value="Estimate Matting", variant="primary"
+                            )
+                            matting_image_reset_btn = gr.Button(value="Reset")
+                        with gr.Accordion("Advanced options", open=False):
+                            image_processing_res = gr.Radio(
+                                [
+                                    ("Native", 0),
+                                    ("Recommended", 768),
+                                ],
+                                label="Processing resolution",
+                                value=default_image_processing_res,
+                            )
+                    with gr.Column():
+                        # dis_image_output_slider = ImageSlider(
+                        #     label="Predicted dichotomous image segmentation",
+                        #     type="filepath",
+                        #     show_download_button=True,
+                        #     show_share_button=True,
+                        #     interactive=False,
+                        #     elem_classes="slider",
+                        #     position=0.25,
+                        # )
+                        matting_image_output = gr.Gallery(label='Output', show_label=False, elem_id="gallery", columns=[1], height='auto')
+                        matting_image_output_files = gr.Files(
+                            label="Matting outputs",
+                            elem_id="download",
+                            interactive=False,
+                        )
+
+                filenames = []
+                filenames.extend(["matting_%d.jpg" %(i+1) for i in range(10)])
+                # example_folder = "images"
+                # print('line 396', __file__)
+                example_folder = os.path.join(os.path.dirname(__file__), "matting_images")
+                # print(example_folder)
+                Examples(
+                    fn=process_pipe_dis,
+                    examples=[
+                        os.path.join(example_folder, name)
+                        for name in filenames
+                    ],
+                    inputs=[dis_image_input],
+                    outputs=[dis_image_output, dis_image_output_files],
+                    # cache_examples=True,
+                    directory_name="images_cache",
+                    cache_examples=False,
+                )
+
 
         ### Image tab
         depth_image_submit_btn.click(
@@ -510,6 +602,38 @@ def run_demo_server(pipe_depth, pipe_normal, pipe_dis):
             queue=False,
         )
 
+        matting_image_submit_btn.click(
+            fn=process_image_check,
+            inputs=matting_image_input,
+            outputs=None,
+            preprocess=False,
+            queue=False,
+        ).success(
+            fn=process_pipe_dis,
+            inputs=[
+                matting_image_input,
+                image_processing_res,
+            ],
+            outputs=[matting_image_output, matting_image_output_files],
+            concurrency_limit=1,
+        )
+
+        matting_image_reset_btn.click(
+            fn=lambda: (
+                None,
+                None,
+                None,
+                default_image_processing_res,
+            ),
+            inputs=[],
+            outputs=[
+                matting_image_input,
+                matting_image_output,
+                matting_image_output_files,
+                image_processing_res,
+            ],
+            queue=False,
+        )
         ### Server launch
 
         demo.queue(
@@ -534,6 +658,7 @@ def main():
                                             use_safetensors=True).to(dtype)
     unet_normal_v1 = UNet2DConditionModel.from_pretrained('guangkaixu/GenPercept', subfolder="unet_normal_v1", use_safetensors=True).to(dtype)
     unet_dis_v1 = UNet2DConditionModel.from_pretrained('guangkaixu/GenPercept', subfolder="unet_dis_v1", use_safetensors=True).to(dtype)
+    unet_matting_v1 = UNet2DConditionModel.from_pretrained('guangkaixu/genpercept-matting', subfolder="unet", use_safetensors=True).to(dtype)
     
     empty_text_embed = torch.from_numpy(np.load("./empty_text_embed.npy")).to(device, dtype)[None] # [1, 77, 1024]
     
@@ -546,19 +671,24 @@ def main():
     pipe_dis = GenPerceptPipeline(vae=vae,
                                   unet=unet_dis_v1,
                                   empty_text_embed=empty_text_embed)
+    pipe_matting = GenPerceptPipeline(vae=vae,
+                                  unet=unet_matting_v1,
+                                  empty_text_embed=empty_text_embed)
     try:
         import xformers
         pipe_depth.enable_xformers_memory_efficient_attention()
         pipe_normal.enable_xformers_memory_efficient_attention()
         pipe_dis.enable_xformers_memory_efficient_attention()
+        pipe_matting.enable_xformers_memory_efficient_attention()
     except:
         pass  # run without xformers
 
     pipe_depth = pipe_depth.to(device)
     pipe_normal = pipe_normal.to(device)
     pipe_dis = pipe_dis.to(device)
+    pipe_matting = pipe_matting.to(device)
     
-    run_demo_server(pipe_depth, pipe_normal, pipe_dis)
+    run_demo_server(pipe_depth, pipe_normal, pipe_dis, pipe_matting)
 
 
 if __name__ == "__main__":