Update app.py

app.py

@@ -9,10 +9,10 @@ import uuid
 
 is_shared_ui = True if "fudan-generative-ai/hallo" in os.environ['SPACE_ID'] else False
 
-if(not is_shared_ui):
+if not is_shared_ui:
     hallo_dir = snapshot_download(repo_id="fudan-generative-ai/hallo", local_dir="pretrained_models")
 
-def run_inference(source_image, driving_audio, progress=gr.Progress(track_tqdm=True)):
+def run_inference(source_image, driving_audio, pose_weight, face_weight, lip_weight, face_expand_ratio, progress=gr.Progress(track_tqdm=True)):
     if is_shared_ui:
         raise gr.Error("This Space only works in duplicated instances")
         
@@ -23,30 +23,44 @@ def run_inference(source_image, driving_audio, progress=gr.Progress(track_tqdm=T
         source_image=source_image,
         driving_audio=driving_audio,
         output=f'output-{unique_id}.mp4',
-        pose_weight=1.5,
-        face_weight=1.0,
-        lip_weight=1.1,
-        face_expand_ratio=1.2,
+        pose_weight=pose_weight,
+        face_weight=face_weight,
+        lip_weight=lip_weight,
+        face_expand_ratio=face_expand_ratio,
         checkpoint=None
     )
     
     inference_process(args)
     return f'output-{unique_id}.mp4' 
 
-
-
+with gr.Blocks() as demo:
+    gr.Markdown(
+        """
+        # Talking Head Generation
+        Upload a face image and driving audio, and adjust the weights to generate a talking head video.
+        """
+    )
+    
     with gr.Row():
         with gr.Column():
-            avatar_face = gr.Image(type="filepath", label="Face")
-            driving_audio = gr.Audio(type="filepath", label="Driving audio")
-            generate = gr.Button("Generate")
+            avatar_face = gr.Image(type="filepath", label="Face", elem_id="face-input")
+            driving_audio = gr.Audio(type="filepath", label="Driving Audio", elem_id="audio-input")
+            
+            with gr.Accordion("Advanced Settings", open=False):
+                pose_weight = gr.Slider(minimum=0.0, value=1.5, label="Pose Weight")
+                face_weight = gr.Slider(minimum=0.0, value=1.0, label="Face Weight")
+                lip_weight = gr.Slider(minimum=0.0, value=1.1, label="Lip Weight")
+                face_expand_ratio = gr.Slider(minimum=0.0, value=1.2, label="Face Expand Ratio")
+            
+            generate = gr.Button("Generate", elem_id="generate-button")
+        
         with gr.Column():
-            output_video = gr.Video(label="Your talking head")
+            output_video = gr.Video(label="Your Talking Head", elem_id="output-video")
 
     generate.click(
         fn=run_inference,
-        inputs=[avatar_face, driving_audio],
+        inputs=[avatar_face, driving_audio, pose_weight, face_weight, lip_weight, face_expand_ratio],
         outputs=output_video
     )
     
-demo.launch()
+demo.launch()