Update app.py

app.py

@@ -126,11 +126,28 @@ gradio_pipeline = GradioPipeline(
 
 # 추가 정의
 elevenlabs_pipeline = ElevenLabsPipeline()
-stf_pipeline = STFPipeline()
+#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):
+
+# @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)
 
 
@@ -161,9 +178,9 @@ def is_square_video(video_path):
     return gr.update(visible=True)
 
 
-def txt_to_driving_video(input_text):
+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_outpath = gpu_wrapped_stf_pipeline_execute(audio_outpath, video_type)
     return video_outpath
 
 
@@ -211,6 +228,7 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
         # 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():
@@ -314,7 +332,7 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
     txt2video_gen_button.click(
         fn=txt_to_driving_video,
         inputs=[
-            script_txt
+            script_txt, video_type
         ],
         outputs=[video_input],
         show_progress=True