Update app.py

app.py

@@ -51,41 +51,52 @@ print("LD_LIBRARY_PATH:", os.environ['LD_LIBRARY_PATH'])
 
 from stf_utils import STFPipeline
 
-stf_pipeline = 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()
-def gpu_wrapped_stf_pipeline_execute(audio_path):
+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)
 
     
 ###### 테스트중 ######
     
 
-stf_pipeline = STFPipeline()
-driving_video_path=gr.Video()
 
 # set tyro theme
 tyro.extras.set_accent_color("bright_cyan")
 args = tyro.cli(ArgumentConfig)
+driving_video_path=gr.Video()
+stf_pipeline_female = STFPipeline()
+stf_pipeline_male = STFPipeline(
+            template_video_path="TEMP/Cam2_2309071202_0012_Natural_Looped.mp4",
+            config_path="TEMP/front_config_v3.json",
+            checkpoint_path="TEMP/0157.pth",
+        )
 
 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")
         input_audio = gr.Audio(label="Upload or Record an audio", type="filepath")
-        stf_button = gr.Button("stf test", variant="primary")
+        stf_button = gr.Button("Audio to Video Generation", variant="primary")
+        video_type = gr.Radio(choices=["Female video", "Male video"], label="Select video type", value="Female video")
         stf_button.click(
-                    fn=gpu_wrapped_stf_pipeline_execute,
-                    inputs=[
-                        input_audio
-                    ],
-                    outputs=[driving_video_path]
-                )
-    with gr.Row():
-        checkbox_female = gr.Checkbox(label="Female video", value=True)  # 디폴트 체크
-        checkbox_male = gr.Checkbox(label="Male video")
+                            fn=gpu_wrapped_stf_pipeline_execute,
+                            inputs=[
+                                input_audio
+                            ],
+                            outputs=[driving_video_path]
+                        )
+        
     with gr.Row():
         driving_video_path.render()