app.py

import os
import gradio as gr
from gradio_client import Client, handle_file
from pathlib import Path
from gradio.utils import get_cache_folder


class Examples(gr.helpers.Examples):
    def __init__(self, *args, cached_folder=None, **kwargs):
        super().__init__(*args, **kwargs, _initiated_directly=False)
        if cached_folder is not None:
            self.cached_folder = cached_folder
            # self.cached_file = Path(self.cached_folder) / "log.csv"
        self.create()


HF_TOKEN = os.environ.get('HF_KEY')

client = Client("Canyu/Diception",
                max_workers=3,
                hf_token=HF_TOKEN)


def process_image_check(path_input):
    if path_input is None:
        raise gr.Error(
            "Missing image in the left pane: please upload an image first."
        )

def infer_image_matting(matting_image_input):
    return client.predict(
      prompt=handle_file(matting_image_input),
      api_name="/test"
    )

def clear_cache():
    return None, None

def run_demo_server():
    gradio_theme = gr.themes.Default()
    with gr.Blocks(
        theme=gradio_theme,
        title="Matting",
    ) as demo:
        with gr.Row():
            gr.Markdown("# Matting Demo")
        with gr.Row():
            gr.Markdown("### Due to the GPU quota limit, if an error occurs, please wait for 5 minutes before retrying.")
        with gr.Row():
            with gr.Column():
                matting_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.Row():
                    img_clear_button = gr.Button("Clear Cache")
                
            with gr.Column():
                # matting_image_output = gr.Image(label='Output')
                matting_image_output =  gr.Image(label='Matting Output')
                    
                        #     label="Matting Output",
                        #     type="filepath",
                        #     show_download_button=True,
                        #     show_share_button=True,
                        #     interactive=False,
                        #     elem_classes="slider",
                        #     position=0.25,
                        # )
                
            

        img_clear_button.click(clear_cache, outputs=[matting_image_input, matting_image_output])

        matting_image_submit_btn.click(
            fn=process_image_check,
            inputs=matting_image_input,
            outputs=None,
            preprocess=False,
            queue=False,
        ).success(
            # fn=process_pipe_matting,
            fn=infer_image_matting,
            inputs=[
                matting_image_input,
            ],
            outputs=[matting_image_output],
            concurrency_limit=1,
        )

        matting_image_reset_btn.click(
            fn=lambda: (
                None,
                None,
            ),
            inputs=[],
            outputs=[
                matting_image_input,
                matting_image_output,
            ],
            queue=False,
        )

        gr.Examples(
            fn=infer_image_matting,
            examples=[
                "assets/person.jpg",              
            ],
            inputs=[matting_image_input],
            outputs=[matting_image_output],
            cache_examples=True,
            # cache_examples=False,
            # cached_folder="cache_dir",
        )
        
    demo.queue(
        api_open=False,
    ).launch()


if __name__ == '__main__':

    run_demo_server()