app.py

# coding: utf-8

"""
The entrance of the gradio
"""

import tyro
import gradio as gr
import os.path as osp
from src.utils.helper import load_description
from src.gradio_pipeline import GradioPipeline
from src.config.crop_config import CropConfig
from src.config.argument_config import ArgumentConfig
from src.config.inference_config import InferenceConfig
import spaces
import cv2
import torch


#추가
from elevenlabs_utils import ElevenLabsPipeline
from setup_environment import initialize_environment
from src.utils.video import extract_audio
#from flux_dev import create_flux_tab
#from flux_schnell import create_flux_tab
# from diffusers import FluxPipeline
import os
import sys

# import gdown
# folder_url = f"https://drive.google.com/drive/folders/1UtKgzKjFAOmZkhNK-OYT0caJ_w2XAnib"
# gdown.download_folder(url=folder_url, output="pretrained_weights", quiet=False)




    
initialize_environment()


sys.path.extend([
    '/home/user/.local/lib/python3.10/site-packages',
    '/home/user/app/',
    '/home/user/app/stf/',
    '/home/user/app/stf/stf_alternative/',
    '/home/user/app/stf/stf_tools'
])

# sys.path.append('/home/user/.local/lib/python3.10/site-packages')
# sys.path.append('/home/user/.local/lib/python3.10/site-packages/stf_alternative/src/stf_alternative')
# sys.path.append('/home/user/.local/lib/python3.10/site-packages/stf_tools/src/stf_tools')
# sys.path.append('/home/user/app/')
# sys.path.append('/home/user/app/stf/')
# sys.path.append('/home/user/app/stf/stf_alternative/')
# sys.path.append('/home/user/app/stf/stf_alternative/src/stf_alternative')
# sys.path.append('/home/user/app/stf/stf_tools')
# sys.path.append('/home/user/app/stf/stf_tools/src/stf_tools')



# CUDA 경로를 환경 변수로 설정
os.environ['PATH'] = '/usr/local/cuda/bin:' + os.environ.get('PATH', '')
os.environ['LD_LIBRARY_PATH'] = '/usr/local/cuda/lib64:' + os.environ.get('LD_LIBRARY_PATH', '')
# 확인용 출력
print("PATH:", os.environ['PATH'])
print("LD_LIBRARY_PATH:", os.environ['LD_LIBRARY_PATH'])

from stf_utils import STFPipeline



# audio_path="assets/examples/driving/test_aud.mp3"
#audio_path_component = gr.Textbox(label="Input", value="assets/examples/driving/test_aud.mp3")

# @spaces.GPU(duration=120)
# def gpu_wrapped_stf_pipeline_execute(audio_path):
#     return stf_pipeline.execute(audio_path)

    
# ###### 테스트중 ######
    

# stf_pipeline = STFPipeline()
# driving_video_path=gr.Video()

# # set tyro theme
# tyro.extras.set_accent_color("bright_cyan")
# args = tyro.cli(ArgumentConfig)

# with gr.Blocks(theme=gr.themes.Soft()) as demo:
#     with gr.Row():
#         audio_path_component = gr.Textbox(label="Input", value="assets/examples/driving/test_aud.mp3")
#         stf_button = gr.Button("stf test", variant="primary")
#         stf_button.click(
#                     fn=gpu_wrapped_stf_pipeline_execute,
#                     inputs=[
#                         audio_path_component
#                     ],
#                     outputs=[driving_video_path]
#                 )
#     with gr.Row():
#         driving_video_path.render()

#     # with gr.Row():
#     #     create_flux_tab()  # image_input을 flux_tab에 전달합니다.

# ###### 테스트중 ######


def partial_fields(target_class, kwargs):
    return target_class(**{k: v for k, v in kwargs.items() if hasattr(target_class, k)})

# set tyro theme
tyro.extras.set_accent_color("bright_cyan")
args = tyro.cli(ArgumentConfig)

# specify configs for inference
inference_cfg = partial_fields(InferenceConfig, args.__dict__)  # use attribute of args to initial InferenceConfig
crop_cfg = partial_fields(CropConfig, args.__dict__)  # use attribute of args to initial CropConfig

gradio_pipeline = GradioPipeline(
    inference_cfg=inference_cfg,
    crop_cfg=crop_cfg,
    args=args
)

# 추가 정의
elevenlabs_pipeline = ElevenLabsPipeline()
#stf_pipeline = STFPipeline()
stf_pipeline_female = STFPipeline()
stf_pipeline_male = STFPipeline(
            template_video_path="TEMP/Cam2_2309071202_0012_Natural_Looped.mp4",
            config_path="front_config_v3.json",
            checkpoint_path="TEMP/0157.pth",
            female_video=False
        )



# @spaces.GPU() #duration=240)
# def gpu_wrapped_stf_pipeline_execute(audio_path):
#     return stf_pipeline.execute(audio_path)

@spaces.GPU()
def gpu_wrapped_stf_pipeline_execute(audio_path, video_type):
    if video_type == "Female video":
        stf_pipeline = stf_pipeline_female
    else:
        stf_pipeline = stf_pipeline_male
        
    return stf_pipeline.execute(audio_path)


@spaces.GPU()
def gpu_wrapped_elevenlabs_pipeline_generate_voice(text, voice):
    return elevenlabs_pipeline.generate_voice(text, voice)



@spaces.GPU()
def gpu_wrapped_execute_video(*args, **kwargs):
    return gradio_pipeline.execute_video(*args, **kwargs)

@spaces.GPU()
def gpu_wrapped_execute_image(*args, **kwargs):
    return gradio_pipeline.execute_image(*args, **kwargs)

def is_square_video(video_path):
    video = cv2.VideoCapture(video_path)

    width = int(video.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(video.get(cv2.CAP_PROP_FRAME_HEIGHT))

    video.release()
    if width != height:
        raise gr.Error("Error: the video does not have a square aspect ratio. We currently only support square videos")

    return gr.update(visible=True)


def txt_to_driving_video(input_text, video_type):
    audio_outpath = gpu_wrapped_elevenlabs_pipeline_generate_voice(text=input_text, voice=None)
    video_outpath = gpu_wrapped_stf_pipeline_execute(audio_outpath, video_type)
    return video_outpath



# assets
title_md = "assets/gradio_title.md"
example_portrait_dir = "assets/examples/source"
example_portrait_dir_custom = "assets/examples/source"
example_video_dir = "assets/examples/driving"
data_examples = [
    [osp.join(example_portrait_dir, "s9.jpg"), osp.join(example_video_dir, "d0.mp4"), True, True, True, True],
    [osp.join(example_portrait_dir, "s6.jpg"), osp.join(example_video_dir, "d0.mp4"), True, True, True, True],
    [osp.join(example_portrait_dir, "s10.jpg"), osp.join(example_video_dir, "d0.mp4"), True, True, True, True],
    [osp.join(example_portrait_dir, "s5.jpg"), osp.join(example_video_dir, "d18.mp4"), True, True, True, True],
    [osp.join(example_portrait_dir, "s7.jpg"), osp.join(example_video_dir, "d19.mp4"), True, True, True, True],
    [osp.join(example_portrait_dir, "s22.jpg"), osp.join(example_video_dir, "d0.mp4"), True, True, True, True],
]
#################### interface logic ####################

# Define components first
eye_retargeting_slider = gr.Slider(minimum=0, maximum=0.8, step=0.01, label="target eyes-open ratio")
lip_retargeting_slider = gr.Slider(minimum=0, maximum=0.8, step=0.01, label="target lip-open ratio")
retargeting_input_image = gr.Image(type="filepath")
output_image = gr.Image(type="numpy")
output_image_paste_back = gr.Image(type="numpy")
output_video = gr.Video()
output_video_concat = gr.Video()

#video_input = gr.Video()
driving_video_path=gr.Video()


with gr.Blocks(theme=gr.themes.Soft()) as demo:
    #gr.HTML(load_description(title_md))

    
    gr.Markdown("# Talk-GEN by ESTsoft")
    gr.Markdown("**Text to talking video generation tool**\n\n")

    #gr.Markdown("### 1. Text to audio")
    gr.Markdown("### 1. Text to Driving-Video")
    with gr.Row():
        
        script_txt = gr.Text()
        # audio_gen_button = gr.Button("Audio generation", variant="primary")
        # with gr.Column():
        #     txt2video_gen_button = gr.Button("txt2video generation", variant="primary")
        video_type = gr.Radio(choices=["Female video", "Male video"], label="Select video type", value="Female video")
        txt2video_gen_button = gr.Button("txt2video generation", variant="primary")

        #with gr.Column():
            #audio_gen_button = gr.Button("Audio generation", variant="primary")
    # with gr.Row():
    #         output_audio = gr.Audio(label="Generated audio", type="filepath")
    
    
    # gr.Markdown("### 2. Audio to Driving-Video")
    # with gr.Row():
    #         #audio_path_component = gr.Textbox(label="Input", value="assets/examples/driving/test_aud.mp3")
    #         video_gen_button = gr.Button("Audio to Video generation", variant="primary")
    # with gr.Row():
    #         #a2v_output = gr.Video()
    #         driving_video_path.render()

        
    gr.Markdown("### 2. Image to Talking-Video with Driving-Video")
    #gr.Markdown(load_description("assets/gradio_description_upload.md"))
    with gr.Row():
        with gr.Accordion(open=True, label="Source Portrait"):
            image_input = gr.Image(type="filepath")
            gr.Examples(
                examples=[
                    #[osp.join(example_portrait_dir, "01.webp")],
                    [osp.join(example_portrait_dir, "02.webp")],
                    [osp.join(example_portrait_dir, "03.jpg")],
                    [osp.join(example_portrait_dir, "04.jpg")],
                    [osp.join(example_portrait_dir, "05.jpg")],
                    [osp.join(example_portrait_dir, "06.jpg")],
                    [osp.join(example_portrait_dir, "07.jpg")],
                    [osp.join(example_portrait_dir, "08.jpg")],
                ],
                inputs=[image_input],
                cache_examples=False,
            )
        with gr.Accordion(open=True, label="Driving Video"):
            video_input = gr.Video()
            gr.Examples(
                examples=[
                    [osp.join(example_video_dir, "d0.mp4")],
                    [osp.join(example_video_dir, "d18.mp4")],
                    [osp.join(example_video_dir, "d19.mp4")],
                    [osp.join(example_video_dir, "d14_trim.mp4")],
                    [osp.join(example_video_dir, "d6_trim.mp4")],
                ],
                inputs=[video_input],
                cache_examples=False,
            )
    with gr.Row():
        with gr.Accordion(open=False, label="Animation Instructions and Options"):
            gr.Markdown(load_description("assets/gradio_description_animation.md"))
            with gr.Row():
                flag_relative_input = gr.Checkbox(value=True, label="relative motion")
                flag_do_crop_input = gr.Checkbox(value=True, label="do crop")
                flag_remap_input = gr.Checkbox(value=True, label="paste-back")
    #gr.Markdown(load_description("assets/gradio_description_animate_clear.md"))
    with gr.Row():
        with gr.Column():
            process_button_animation = gr.Button("🚀 Animate", variant="primary")
        with gr.Column():
            process_button_reset = gr.ClearButton([image_input, video_input, output_video, output_video_concat], value="🧹 Clear")
    with gr.Row():
        with gr.Column():
            with gr.Accordion(open=True, label="The animated video in the original image space"):
                output_video.render()
        with gr.Column():
            with gr.Accordion(open=True, label="The animated video"):
                output_video_concat.render()
    # with gr.Row():
    #     # Examples
    #     gr.Markdown("## You could also choose the examples below by one click ⬇️")
    # with gr.Row():
    #     gr.Examples(
    #         examples=data_examples,
    #         fn=gpu_wrapped_execute_video,
    #         inputs=[
    #             image_input,
    #             video_input,
    #             flag_relative_input,
    #             flag_do_crop_input,
    #             flag_remap_input
    #         ],
    #         outputs=[output_image, output_image_paste_back],
    #         examples_per_page=6,
    #         cache_examples=False,
    #     )

    process_button_animation.click(
        fn=gpu_wrapped_execute_video,
        inputs=[
            image_input,
            video_input,
            flag_relative_input,
            flag_do_crop_input,
            flag_remap_input
        ],
        outputs=[output_video, output_video_concat],
        show_progress=True
    )
    txt2video_gen_button.click(
        fn=txt_to_driving_video,
        inputs=[
            script_txt, video_type
        ],
        outputs=[video_input],
        show_progress=True
    )
    # audio_gen_button.click(
    #     fn=gpu_wrapped_elevenlabs_pipeline_generate_voice,
    #     inputs=[
    #         script_txt
    #     ],
    #     outputs=[output_audio],
    #     show_progress=True
    # )

    # video_gen_button.click(
    #     fn=gpu_wrapped_stf_pipeline_execute,
    #     inputs=[
    #         output_audio
    #         #audio_path_component
    #     ],
    #     outputs=[
    #         video_input
    #         #driving_video_path
    #     ],
    #     show_progress=True
    # )


    
    # image_input.change(
    #     fn=gradio_pipeline.prepare_retargeting,
    #     inputs=image_input,
    #     outputs=[eye_retargeting_slider, lip_retargeting_slider, retargeting_input_image]
    # )
    video_input.upload(
        fn=is_square_video,
        inputs=video_input,
        outputs=video_input
    )


demo.launch(
    server_port=args.server_port,
    share=args.share,
    server_name=args.server_name
)