Update app.py

app.py

@@ -22,7 +22,7 @@ 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 flux_schnell import create_flux_tab
 # from diffusers import FluxPipeline
 
 # import gdown
@@ -31,73 +31,7 @@ from src.utils.video import extract_audio
 
 
 
-# #========================= # FLUX 모델 로드 설정
-# flux_pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=torch.bfloat16)
-# flux_pipe.enable_sequential_cpu_offload()
-# flux_pipe.vae.enable_slicing()
-# flux_pipe.vae.enable_tiling()
-# flux_pipe.to(torch.float16)
 
-
-# @spaces.GPU(duration=120)
-# def generate_image(prompt, guidance_scale, width, height):
-#     # 이미지를 생성하는 함수
-#     output_image = flux_pipe(
-#         prompt=prompt,
-#         guidance_scale=guidance_scale,
-#         height=height,
-#         width=width,
-#         num_inference_steps=4,
-#         max_sequence_length=256,
-#     ).images[0]
-
-#     # 결과 폴더 생성
-#     result_folder = "/tmp/flux/"
-#     os.makedirs(result_folder, exist_ok=True)
-
-#     # 파일 이름 생성
-#     timestamp = datetime.now().strftime("%Y%m%d%H%M%S")
-#     #filename = f"{prompt.replace(' ', '_')}_{timestamp}.png"
-#     filename = f"{'_'.join(prompt.split()[:3])}_{timestamp}.png"
-#     output_path = os.path.join(result_folder, filename)
-
-#     # # 이미지를 저장
-#     # output_image.save(output_path)
-
-#     return output_image, output_path  # 두 개의 출력 반환
-
-# def flux_tab(): #image_input):  # image_input을 인자로 받습니다.
-#     with gr.Tab("FLUX 이미지 생성"):
-#         with gr.Row():
-#             with gr.Column():
-#                 # 사용자 입력 설정
-#                 prompt = gr.Textbox(label="Prompt", value="A cat holding a sign that says hello world")
-#                 guidance_scale = gr.Slider(label="Guidance Scale", minimum=0.0, maximum=20.0, value=3.5, step=0.1)
-#                 width = gr.Slider(label="Width", minimum=256, maximum=2048, value=512, step=64)
-#                 height = gr.Slider(label="Height", minimum=256, maximum=2048, value=512, step=64)
-
-#             with gr.Column():
-#                 # 출력 이미지와 다운로드 버튼
-#                 output_image = gr.Image(type="pil", label="Output")
-#                 download_button = gr.File(label="Download")
-#                 generate_button = gr.Button("이미지 생성")
-#                 #use_in_text2lipsync_button = gr.Button("이 이미지를 Text2Lipsync에서 사용하기")  # 새로운 버튼 추가
-
-#                 # 클릭 이벤트를 정의
-#                 generate_button.click(
-#                     fn=generate_image,
-#                     inputs=[prompt, guidance_scale, width, height],
-#                     outputs=[output_image, download_button]
-#                 )
-
-#                 # # 새로운 버튼 클릭 이벤트 정의
-#                 # use_in_text2lipsync_button.click(
-#                 #     fn=lambda img: img,  # 간단한 람다 함수를 사용하여 이미지를 그대로 전달
-#                 #     inputs=[output_image],  # 생성된 이미지를 입력으로 사용
-#                 #     outputs=[image_input]  # Text to LipSync 탭의 image_input을 업데이트
-#                 # )
-
-# #========================= # FLUX 모델 로드 설정
     
 initialize_environment()
 
@@ -128,272 +62,271 @@ 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)
+# @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()
+# stf_pipeline = STFPipeline()
+# driving_video_path=gr.Video()
 
-# set tyro theme
-tyro.extras.set_accent_color("bright_cyan")
-args = tyro.cli(ArgumentConfig)
+# # 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.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에 전달합니다.
+#     # 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)})
+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)
+# 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
+# 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
-# )
+gradio_pipeline = GradioPipeline(
+    inference_cfg=inference_cfg,
+    crop_cfg=crop_cfg,
+    args=args
+)
 
-# # 추가 정의
-# elevenlabs_pipeline = ElevenLabsPipeline()
-# stf_pipeline = STFPipeline()
+# 추가 정의
+elevenlabs_pipeline = ElevenLabsPipeline()
+stf_pipeline = STFPipeline()
 
 
-# @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_stf_pipeline_execute(audio_path):
+    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_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_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)
+@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)
+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))
+    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")
+    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)
+    return gr.update(visible=True)
 
-# def txt_to_driving_video(text):
-#     audio_path = gpu_wrapped_elevenlabs_pipeline_generate_voice(text)
-#     driving_video_path = gpu_wrapped_stf_pipeline_execute(audio_path)
-#     return driving_video_path
 
-# # 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()
+# 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 ####################
 
-# # video_input = gr.Video()
-# driving_video_path=gr.Video()
+# 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))
-
-#     with gr.Tabs():
-#         with gr.Tab("Text to LipSync"):
-#             gr.Markdown("# Text to LipSync")
-#             with gr.Row():
-#                 with gr.Column():
-#                     script_txt = gr.Text()
-#                 # with gr.Column():
-#                 #     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")
-#             with gr.Row():
-#                     video_gen_button = gr.Button("Audio to Video generation", variant="primary")
+
+with gr.Blocks(theme=gr.themes.Soft()) as demo:
+    #gr.HTML(load_description(title_md))
+
+    with gr.Tabs():
+        with gr.Tab("Text to LipSync"):
+            gr.Markdown("# Text to LipSync")
+            with gr.Row():
+                with gr.Column():
+                    script_txt = gr.Text()
+                # with gr.Column():
+                #     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")
+            with gr.Row():
+                    video_gen_button = gr.Button("Audio to Video generation", variant="primary")
+            with gr.Row():
+                    a2v_output = gr.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,
-#                 )
+            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
-#             #     ],
-#             #     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
-#                 ],
-#                 outputs=[video_input],
-#                 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
+            )
+            # txt2video_gen_button.click(
+            #     fn=txt_to_driving_video,
+            #     inputs=[
+            #         script_txt
+            #     ],
+            #     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
+                ],
+                outputs=[a2v_output],
+                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
-#             )
+            # 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
+            )
         
-#         # 세 번째 탭: Flux 개발용 탭
-#         with gr.Tab("FLUX Image"):
-#             flux_demo = create_flux_tab(image_input)  # Flux 개발용 탭 생성
+        # 세 번째 탭: Flux 개발용 탭
+        with gr.Tab("FLUX Image"):
+            flux_demo = create_flux_tab(image_input)  # Flux 개발용 탭 생성
 
 demo.launch(
     server_port=args.server_port,