KlingTeam/LivePortrait

1# coding: utf-82 3"""4The entrance of the gradio5"""6 7import tyro8import gradio as gr9import os.path as osp10from src.utils.helper import load_description11from src.gradio_pipeline import GradioPipeline12from src.config.crop_config import CropConfig13from src.config.argument_config import ArgumentConfig14from src.config.inference_config import InferenceConfig15import spaces16import cv217 18# import gdown19# folder_url = f"https://drive.google.com/drive/folders/1UtKgzKjFAOmZkhNK-OYT0caJ_w2XAnib"20# gdown.download_folder(url=folder_url, output="pretrained_weights", quiet=False)21 22def partial_fields(target_class, kwargs):23    return target_class(**{k: v for k, v in kwargs.items() if hasattr(target_class, k)})24 25# set tyro theme26tyro.extras.set_accent_color("bright_cyan")27args = tyro.cli(ArgumentConfig)28 29# specify configs for inference30inference_cfg = partial_fields(InferenceConfig, args.__dict__)  # use attribute of args to initial InferenceConfig31crop_cfg = partial_fields(CropConfig, args.__dict__)  # use attribute of args to initial CropConfig32 33gradio_pipeline = GradioPipeline(34    inference_cfg=inference_cfg,35    crop_cfg=crop_cfg,36    args=args37)38 39@spaces.GPU(duration=240)40def gpu_wrapped_execute_video(*args, **kwargs):41    return gradio_pipeline.execute_video(*args, **kwargs)42 43@spaces.GPU(duration=240)44def gpu_wrapped_execute_image(*args, **kwargs):45    return gradio_pipeline.execute_image(*args, **kwargs)46 47def is_square_video(video_path):48    video = cv2.VideoCapture(video_path)49 50    width = int(video.get(cv2.CAP_PROP_FRAME_WIDTH))51    height = int(video.get(cv2.CAP_PROP_FRAME_HEIGHT))52 53    video.release()54    if width != height:55        raise gr.Error("Error: the video does not have a square aspect ratio. We currently only support square videos")56 57    return gr.update(visible=True)58 59# assets60title_md = "assets/gradio_title.md"61example_portrait_dir = "assets/examples/source"62example_video_dir = "assets/examples/driving"63data_examples = [64    [osp.join(example_portrait_dir, "s9.jpg"), osp.join(example_video_dir, "d0.mp4"), True, True, True, True],65    [osp.join(example_portrait_dir, "s6.jpg"), osp.join(example_video_dir, "d0.mp4"), True, True, True, True],66    [osp.join(example_portrait_dir, "s10.jpg"), osp.join(example_video_dir, "d0.mp4"), True, True, True, True],67    [osp.join(example_portrait_dir, "s5.jpg"), osp.join(example_video_dir, "d18.mp4"), True, True, True, True],68    [osp.join(example_portrait_dir, "s7.jpg"), osp.join(example_video_dir, "d19.mp4"), True, True, True, True],69    [osp.join(example_portrait_dir, "s22.jpg"), osp.join(example_video_dir, "d0.mp4"), True, True, True, True],70]71#################### interface logic ####################72 73# Define components first74eye_retargeting_slider = gr.Slider(minimum=0, maximum=0.8, step=0.01, label="target eyes-open ratio")75lip_retargeting_slider = gr.Slider(minimum=0, maximum=0.8, step=0.01, label="target lip-open ratio")76retargeting_input_image = gr.Image(type="filepath")77output_image = gr.Image(type="numpy")78output_image_paste_back = gr.Image(type="numpy")79output_video = gr.Video()80output_video_concat = gr.Video()81 82with gr.Blocks(theme=gr.themes.Soft()) as demo:83    gr.HTML(load_description(title_md))84    gr.Markdown(load_description("assets/gradio_description_upload.md"))85    with gr.Row():86        with gr.Accordion(open=True, label="Source Portrait"):87            image_input = gr.Image(type="filepath")88            gr.Examples(89                examples=[90                    [osp.join(example_portrait_dir, "s9.jpg")],91                    [osp.join(example_portrait_dir, "s6.jpg")],92                    [osp.join(example_portrait_dir, "s10.jpg")],93                    [osp.join(example_portrait_dir, "s5.jpg")],94                    [osp.join(example_portrait_dir, "s7.jpg")],95                    [osp.join(example_portrait_dir, "s12.jpg")],96                    [osp.join(example_portrait_dir, "s22.jpg")],97                ],98                inputs=[image_input],99                cache_examples=False,100            )101        with gr.Accordion(open=True, label="Driving Video"):102            video_input = gr.Video()103            gr.Examples(104                examples=[105                    [osp.join(example_video_dir, "d0.mp4")],106                    [osp.join(example_video_dir, "d18.mp4")],107                    [osp.join(example_video_dir, "d19.mp4")],108                    [osp.join(example_video_dir, "d14_trim.mp4")],109                    [osp.join(example_video_dir, "d6_trim.mp4")],110                ],111                inputs=[video_input],112                cache_examples=False,113            )114    with gr.Row():115        with gr.Accordion(open=False, label="Animation Instructions and Options"):116            gr.Markdown(load_description("assets/gradio_description_animation.md"))117            with gr.Row():118                flag_relative_input = gr.Checkbox(value=True, label="relative motion")119                flag_do_crop_input = gr.Checkbox(value=True, label="do crop")120                flag_remap_input = gr.Checkbox(value=True, label="paste-back")121    gr.Markdown(load_description("assets/gradio_description_animate_clear.md"))122    with gr.Row():123        with gr.Column():124            process_button_animation = gr.Button("🚀 Animate", variant="primary")125        with gr.Column():126            process_button_reset = gr.ClearButton([image_input, video_input, output_video, output_video_concat], value="🧹 Clear")127    with gr.Row():128        with gr.Column():129            with gr.Accordion(open=True, label="The animated video in the original image space"):130                output_video.render()131        with gr.Column():132            with gr.Accordion(open=True, label="The animated video"):133                output_video_concat.render()134    with gr.Row():135        # Examples136        gr.Markdown("## You could also choose the examples below by one click ⬇️")137    with gr.Row():138        gr.Examples(139            examples=data_examples,140            fn=gpu_wrapped_execute_video,141            inputs=[142                image_input,143                video_input,144                flag_relative_input,145                flag_do_crop_input,146                flag_remap_input147            ],148            outputs=[output_image, output_image_paste_back],149            examples_per_page=6,150            cache_examples=False,151        )152    gr.Markdown(load_description("assets/gradio_description_retargeting.md"), visible=True)153    with gr.Row(visible=True):154        eye_retargeting_slider.render()155        lip_retargeting_slider.render()156    with gr.Row(visible=True):157        process_button_retargeting = gr.Button("🚗 Retargeting", variant="primary")158        process_button_reset_retargeting = gr.ClearButton(159            [160                eye_retargeting_slider,161                lip_retargeting_slider,162                retargeting_input_image,163                output_image,164                output_image_paste_back165            ],166            value="🧹 Clear"167        )168    with gr.Row(visible=True):169        with gr.Column():170            with gr.Accordion(open=True, label="Retargeting Input"):171                retargeting_input_image.render()172                gr.Examples(173                    examples=[174                        [osp.join(example_portrait_dir, "s9.jpg")],175                        [osp.join(example_portrait_dir, "s6.jpg")],176                        [osp.join(example_portrait_dir, "s10.jpg")],177                        [osp.join(example_portrait_dir, "s5.jpg")],178                        [osp.join(example_portrait_dir, "s7.jpg")],179                        [osp.join(example_portrait_dir, "s12.jpg")],180                        [osp.join(example_portrait_dir, "s22.jpg")],181                    ],182                    inputs=[retargeting_input_image],183                    cache_examples=False,184                )185        with gr.Column():186            with gr.Accordion(open=True, label="Retargeting Result"):187                output_image.render()188        with gr.Column():189            with gr.Accordion(open=True, label="Paste-back Result"):190                output_image_paste_back.render()191    # binding functions for buttons192    process_button_retargeting.click(193        # fn=gradio_pipeline.execute_image,194        fn=gpu_wrapped_execute_image,195        inputs=[eye_retargeting_slider, lip_retargeting_slider, retargeting_input_image, flag_do_crop_input],196        outputs=[output_image, output_image_paste_back],197        show_progress=True198    )199    process_button_animation.click(200        fn=gpu_wrapped_execute_video,201        inputs=[202            image_input,203            video_input,204            flag_relative_input,205            flag_do_crop_input,206            flag_remap_input207        ],208        outputs=[output_video, output_video_concat],209        show_progress=True210    )211    # image_input.change(212    #     fn=gradio_pipeline.prepare_retargeting,213    #     inputs=image_input,214    #     outputs=[eye_retargeting_slider, lip_retargeting_slider, retargeting_input_image]215    # )216    video_input.upload(217        fn=is_square_video,218        inputs=video_input,219        outputs=video_input220    )221 222demo.launch(223    server_port=args.server_port,224    share=args.share,225    server_name=args.server_name226)227