Update app.py

app.py

@@ -14,7 +14,11 @@ hf_hub_download(
     local_dir="checkpoint"
 )
 
-def extract_frames_with_labels(video_path, output_dir="frames"):
+def extract_frames_with_labels(video_path, base_output_dir="frames"):
+    # Generate a timestamped folder name
+    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+    output_dir = os.path.join(base_output_dir, f"frames_{timestamp}")
+    
     # Ensure output directory exists
     os.makedirs(output_dir, exist_ok=True)
     
@@ -51,27 +55,19 @@ def extract_frames_with_labels(video_path, output_dir="frames"):
     return frame_data
 
 # Define a function to run your script with selected inputs
-def run_xportrait(
-    model_config, 
-    output_dir_base, 
-    resume_dir, 
-    seed, 
-    uc_scale, 
-    source_image, 
-    driving_video, 
-    best_frame, 
-    out_frames, 
-    num_mix, 
-    ddim_steps
-):
+def run_xportrait(source_image, driving_video, seed, uc_scale, best_frame, out_frames, num_mix, ddim_steps):
     # Check if the model weights are in place
     if not os.path.exists(resume_dir):
         return "Model weights not found in checkpoint directory. Please download them first."
 
     # Create a unique output directory name based on current date and time
+    output_dir_base = "outputs"
     timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
     output_dir = os.path.join(output_dir_base, f"output_{timestamp}")
     os.makedirs(output_dir, exist_ok=True)
+
+    model_config = "config/cldm_v15_appearance_pose_local_mm.yaml"
+    resume_dir = "checkpoint/model_state-415001.th"
     
     # Construct the command
     command = [
@@ -94,7 +90,7 @@ def run_xportrait(
         subprocess.run(command, check=True)
         
         # Find the generated video file in the output directory
-        video_files = glob.glob(os.path.join(output_dir, "*.mp4")) + glob.glob(os.path.join(output_dir, "*.avi"))
+        video_files = glob.glob(os.path.join(output_dir, "*.mp4"))
         print(video_files)
         if video_files:
             return f"Output video saved at: {video_files[0]}", video_files[0]
@@ -104,25 +100,41 @@ def run_xportrait(
         return f"An error occurred: {e}", None
 
 # Set up Gradio interface
-app = gr.Interface(
-    fn=run_xportrait,
-    inputs=[
-        gr.Textbox(value="config/cldm_v15_appearance_pose_local_mm.yaml", label="Model Config Path"),
-        gr.Textbox(value="outputs", label="Output Directory"),
-        gr.Textbox(value="checkpoint/model_state-415001.th", label="Resume Directory"),
-        gr.Number(value=999, label="Seed"),
-        gr.Number(value=5, label="UC Scale"),
-        gr.Image(label="Source Image", type="filepath"),
-        gr.Video(label="Driving Video"),
-        gr.Number(value=36, label="Best Frame"),
-        gr.Number(value=-1, label="Out Frames"),
-        gr.Number(value=4, label="Number of Mix"),
-        gr.Number(value=30, label="DDIM Steps")
-    ],
-    outputs=["text", "video"],
-    title="XPortrait Model Runner",
-    description="Run XPortrait with customizable parameters."
-)
+with gr.Blocks() as demo:
+    with gr.Column(elem_id="col-container"):
+        with gr.Row():
+            with gr.Column():
+                source_image = gr.Image(label="Source Image", type="filepath")
+                driving_video = gr.Video(label="Driving Video")
+                with gr.Row():
+                    seed = gr.Number(value=999, label="Seed")
+                    uc_scale = gr.Number(value=5, label="UC Scale")
+                with gr.Group():
+                    with gr.Row():
+                        best_frame = gr.Number(value=36, label="Best Frame")
+                        out_frames = gr.Number(value=-1, label="Out Frames")
+                    with gr.Accordion("Driving video Frames"):
+                        driving_frames = gr.Gallery(show_label=True)
+                with gr.Row():
+                    num_mix = gr.Number(value=4, label="Number of Mix")
+                    ddim_steps = gr.Number(value=30, label="DDIM Steps")
+                submit_btn = gr.Button("Submit")
+            with gr.Column():
+                video_output = gr.Video(label="Output Video")
+                status = gr.Textbox(label="status")
+
+    driving_video.upload(
+        fn = extract_frames_with_labels,
+        inputs = [driving_video],
+        ouputs = [driving_frames],
+        queue = False
+    )
+    
+    submit_btn.click(
+        fn = run_xportrait,
+        inputs = [source_image, driving_video, seed, uc_scale, best_frame, out_frames, num_mix, ddim_steps],
+        outputs = [status, video_output]
+    )
 
 # Launch the Gradio app
-app.launch()
+demo.launch()