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 diffusers import FluxPipeline

# 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)



#========================= # 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()

import sys
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')


import os
# 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():
        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()

# @spaces.GPU(duration=200)
# def gpu_wrapped_elevenlabs_pipeline_generate_voice(text, voice):
#     return elevenlabs_pipeline.generate_voice(text, voice)


    

# @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)

# 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()

# 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():
#                     audio_gen_button = gr.Button("Audio generation", variant="primary")
#             with gr.Row():
#                     output_audio_path = gr.Audio(label="Generated audio", type="filepath")
                
#             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,
#                 )
        
#             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
#             )
#             audio_gen_button.click(
#                 fn=gpu_wrapped_elevenlabs_pipeline_generate_voice,
#                 inputs=[
#                     script_txt
#                 ],
#                 outputs=[output_audio_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
#             )
        
#         # 세 번째 탭: Flux 개발용 탭
#         with gr.Tab("FLUX Dev"):
#             flux_demo = create_flux_tab(image_input)  # Flux 개발용 탭 생성

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