import gradio as gr
from FONT.gradio_demo import FONT
from Demo_TFR_Pirenderer.gradio_demo import OPT

def tfg_FONT(source_img, driving_audio):
    tfg = FONT()
    video_path = tfg.test(source_image_path=source_img, driving_audio_path=driving_audio)
    return video_path

tfg_FONT_demo = gr.Interface(fn=tfg_FONT, inputs=[gr.Image(type="filepath"), gr.Audio(type="filepath")],
                        outputs=gr.Video(include_audio=True, height=256, width=256),
                    examples=[["./example/images/60.png", "./example/audios/6343252661930009508_00092.wav"],
                              ["./example/images/7.png", "./example/audios/6350921755403330389_00056.wav"]])

def tfg_OPT(source_img, driving_audio):
    tfg = OPT(checkpoint_path='./Demo_TFR_Pirenderer/checkpoints', config_path='./Demo_TFR_Pirenderer/src/config')
    video_path = tfg.test(source_image=source_img, driven_audio=driving_audio)
    return video_path

tfr_OPT_demo = gr.Interface(fn=tfg_OPT, inputs=[gr.Image(type="filepath"), gr.Audio(type="filepath")],
                        outputs=gr.Video(include_audio=True, height=256, width=256),
                    examples=[["./example/images/60.png", "./example/audios/6343252661930009508_00092.wav"],
                              ["./example/images/7.png", "./example/audios/6350921755403330389_00056.wav"]])


demo = gr.TabbedInterface([tfr_OPT_demo, tfg_FONT_demo], ["PIRenderer_based_Talking_Face_Generation", "FOMM_based_Talking_Face_Generation"])


info = demo.launch()
print(info)