Update app.py

app.py

@@ -20,13 +20,81 @@ import cv2
 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_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)
+
+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 = "result_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
@@ -86,6 +154,9 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
     with gr.Row():
         driving_video_path.render()
 
+    with gr.Row():
+        flux_tab()  # image_input을 flux_tab에 전달합니다.
+