Create app.py

app.py

@@ -0,0 +1,391 @@
+# 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 gdown
+# folder_url = f"https://drive.google.com/drive/folders/1UtKgzKjFAOmZkhNK-OYT0caJ_w2XAnib"
+# gdown.download_folder(url=folder_url, output="pretrained_weights", quiet=False)
+
+import sys
+from src.utils.video import extract_audio
+from elevenlabs_utils import ElevenLabsPipeline
+from setup_environment import initialize_environment
+
+initialize_environment()
+
+
+
+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
+
+
+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
+)
+
+@spaces.GPU(duration=240)
+def gpu_wrapped_execute_video(*args, **kwargs):
+    return gradio_pipeline.execute_video(*args, **kwargs)
+
+@spaces.GPU(duration=240)
+def gpu_wrapped_execute_image(*args, **kwargs):
+    return gradio_pipeline.execute_image(*args, **kwargs)
+
+@spaces.GPU(duration=240)
+def gpu_wrapped_stf_pipeline_execute(audio_path):
+    return stf_pipeline.execute(audio_path)
+
+@spaces.GPU(duration=240)
+def gpu_wrapped_elevenlabs_pipeline_generate_voice(text, voice):
+    return elevenlabs_pipeline.generate_voice(text, voice)
+
+    
+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)
+    
+
+# assets
+title_md = "assets/gradio_title.md"
+example_portrait_dir = "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()
+
+        
+
+
+def run_end_to_end(image_path, text, voice, input_video, flag_relative, flag_do_crop, flag_remap, flag_crop_driving_video, male, animal):
+    
+    # # animal 체크 여부에 따라 다른 pipeline 사용
+    # if animal:
+    #     gradio_pipeline = GradioPipelineAnimal(
+    #         inference_cfg=inference_cfg,
+    #         crop_cfg=crop_cfg,
+    #         args=args
+    #     )
+    # else:
+    #     gradio_pipeline = GradioPipeline(
+    #         inference_cfg=inference_cfg,
+    #         crop_cfg=crop_cfg,
+    #         args=args
+    #     )
+
+    
+    if not male:
+        stf_pipeline = STFPipeline()
+    else:
+        stf_pipeline = STFPipeline(template_video_path="/home/user/app/stf/TEMP/Cam2_2309071202_0012_Natural_Looped.mp4",
+           config_path="/home/user/app/stf/TEMP/front_config_v3.json",
+           checkpoint_path="/home/user/app/stf/TEMP/0157.pth",
+           )
+
+    if input_video is None:
+        #audio_path = elevenlabs_pipeline.generate_voice(text, voice)
+        audio_path = gpu_wrapped_elevenlabs_pipeline_generate_voice(text, voice)
+        #driving_video_path = stf_pipeline.execute(audio_path)
+        driving_video_path = gpu_wrapped_stf_pipeline_execute(audio_path)
+    else:
+        driving_video_path = input_video
+        audio_path = osp.join("animations", Path(input_video).stem+".wav")
+        extract_audio(driving_video_path, audio_path)
+
+
+    #output_path, crop_output_path = gradio_pipeline.execute_video(
+    output_path, crop_output_path = gpu_wrapped_execute_video(
+            input_source_image_path=image_path,
+            input_driving_video_path=driving_video_path,
+            # input_driving_video_pickle_path=None,
+            flag_do_crop_input=flag_do_crop,
+            # flag_remap_input=flag_remap,
+            # driving_multiplier=1.0,
+            # flag_stitching=False,
+            flag_crop_driving_video_input=flag_crop_driving_video,
+            # scale=2.3,
+            # vx_ratio=0.0,
+            # vy_ratio=-0.125,
+            # scale_crop_driving_video=2.2,
+            # vx_ratio_crop_driving_video=0.0,
+            # vy_ratio_crop_driving_video=-0.1,
+            # tab_selection=None,
+            audio_path=audio_path
+            )
+   
+    return output_path, crop_output_path
+
+
+
+with gr.Blocks(theme=gr.themes.Soft()) as demo:
+    with gr.Tabs():
+        # 첫 번째 탭: Text to LipSync
+        with gr.Tab("Text to LipSync"):
+            gr.Markdown("# Text to LipSync")
+            with gr.Row():
+                script_txt = gr.Text()
+                voice = gr.Audio(label="사용자 음성", type="filepath")
+                input_video = gr.Video()
+
+            with gr.Row():
+                image_input = gr.Image(type="filepath")  # 여기서 image_input을 정의합니다.
+                output_video.render()
+                crop_output_video.render()
+
+            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")
+                flag_crop_driving_video_input = gr.Checkbox(value=False, label="do crop (driving video)")
+                male = gr.Checkbox(value=False, label="male")
+                animal = gr.Checkbox(value=False, label="animal")  # animal 체크박스 추가
+
+            with gr.Row():
+                generate_speech = gr.Button("🚀 Generate Speech", variant="primary")
+
+            generate_speech.click(
+                fn=run_end_to_end,
+                inputs=[
+                    image_input,
+                    script_txt,
+                    voice,
+                    input_video,
+                    flag_relative_input,
+                    flag_do_crop_input,
+                    flag_remap_input,
+                    flag_crop_driving_video_input,
+                    male,
+                    animal  # 추가된 animal 입력
+                ],
+                outputs=[output_video, crop_output_video]
+            )
+
+        # # 두 번째 탭: FLUX 이미지 생성
+        # with gr.Tab("FLUX 이미지 생성"):
+        #     flux_tab(image_input)  # FLUX 이미지 생성을 위한 별도의 탭
+
+        # # 세 번째 탭: Flux 개발용 탭
+        # with gr.Tab("FLUX Dev"):
+        #     flux_demo = create_flux_tab()  # Flux 개발용 탭 생성
+        #     #flux_demo.render()  # 해당 UI를 별도의 탭에만 렌더링
+
+
+
+
+
+
+
+    
+# with gr.Blocks(theme=gr.themes.Soft()) as demo:
+#     gr.HTML(load_description(title_md))
+#     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, "s9.jpg")],
+#                     [osp.join(example_portrait_dir, "s6.jpg")],
+#                     [osp.join(example_portrait_dir, "s10.jpg")],
+#                     [osp.join(example_portrait_dir, "s5.jpg")],
+#                     [osp.join(example_portrait_dir, "s7.jpg")],
+#                     [osp.join(example_portrait_dir, "s12.jpg")],
+#                     [osp.join(example_portrait_dir, "s22.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,
+#         )
+#     gr.Markdown(load_description("assets/gradio_description_retargeting.md"), visible=True)
+#     with gr.Row(visible=True):
+#         eye_retargeting_slider.render()
+#         lip_retargeting_slider.render()
+#     with gr.Row(visible=True):
+#         process_button_retargeting = gr.Button("🚗 Retargeting", variant="primary")
+#         process_button_reset_retargeting = gr.ClearButton(
+#             [
+#                 eye_retargeting_slider,
+#                 lip_retargeting_slider,
+#                 retargeting_input_image,
+#                 output_image,
+#                 output_image_paste_back
+#             ],
+#             value="🧹 Clear"
+#         )
+#     with gr.Row(visible=True):
+#         with gr.Column():
+#             with gr.Accordion(open=True, label="Retargeting Input"):
+#                 retargeting_input_image.render()
+#                 gr.Examples(
+#                     examples=[
+#                         [osp.join(example_portrait_dir, "s9.jpg")],
+#                         [osp.join(example_portrait_dir, "s6.jpg")],
+#                         [osp.join(example_portrait_dir, "s10.jpg")],
+#                         [osp.join(example_portrait_dir, "s5.jpg")],
+#                         [osp.join(example_portrait_dir, "s7.jpg")],
+#                         [osp.join(example_portrait_dir, "s12.jpg")],
+#                         [osp.join(example_portrait_dir, "s22.jpg")],
+#                     ],
+#                     inputs=[retargeting_input_image],
+#                     cache_examples=False,
+#                 )
+#         with gr.Column():
+#             with gr.Accordion(open=True, label="Retargeting Result"):
+#                 output_image.render()
+#         with gr.Column():
+#             with gr.Accordion(open=True, label="Paste-back Result"):
+#                 output_image_paste_back.render()
+#     # binding functions for buttons
+#     process_button_retargeting.click(
+#         # fn=gradio_pipeline.execute_image,
+#         fn=gpu_wrapped_execute_image,
+#         inputs=[eye_retargeting_slider, lip_retargeting_slider, retargeting_input_image, flag_do_crop_input],
+#         outputs=[output_image, output_image_paste_back],
+#         show_progress=True
+#     )
+#     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
+#     )
+#     # 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
+)