Update app.py

app.py

@@ -6,6 +6,8 @@ from huggingface_hub import hf_hub_download
 import glob
 from datetime import datetime
 
+is_shared_ui = True if "fffiloni/X-Portrait" in os.environ['SPACE_ID'] else False
+
 # Ensure 'checkpoint' directory exists
 os.makedirs("checkpoint", exist_ok=True)
 
@@ -15,7 +17,31 @@ hf_hub_download(
     local_dir="checkpoint"
 )
 
+def trim_video(video_path, output_dir="trimmed_videos", max_duration=2):
+    # Create output directory if it does not exist
+    os.makedirs(output_dir, exist_ok=True)
+    
+    # Generate a timestamp for the output filename
+    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+    output_path = os.path.join(output_dir, f"trimmed_video_{timestamp}.mp4")
+    
+    # Load the video
+    with VideoFileClip(video_path) as video:
+        # Check the duration of the video
+        if video.duration > max_duration:
+            # Trim the video to the first max_duration seconds
+            trimmed_video = video.subclip(0, max_duration)
+            # Write the trimmed video to a file
+            trimmed_video.write_videofile(output_path, codec="libx264")
+            return output_path
+        else:
+            # If the video is within the duration, return the original path
+            return video_path
+
 def extract_frames_with_labels(video_path, base_output_dir="frames"):
+    if is_shared_ui :
+        video_path = trim_video(video_path)
+        print("Path to the (trimmed) driving video:", video_path)
     # 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}")
@@ -106,6 +132,7 @@ div#frames-gallery{
 with gr.Blocks(css=css) as demo:
     with gr.Column(elem_id="col-container"):
         gr.Markdown("# X-Portrait: Expressive Portrait Animation with Hierarchical Motion Attention")
+        gr.Markdown("On this shared UI, drinving video input will be trimmed to 2 seconds max. Duplicate this space for more controls.")
         gr.HTML("""
         <div style="display:flex;column-gap:4px;">
             <a href='https://github.com/bytedance/X-Portrait'>
@@ -123,8 +150,8 @@ with gr.Blocks(css=css) as demo:
                     driving_video = gr.Video(label="Driving Video")
                 with gr.Group():
                     with gr.Row():
-                        best_frame = gr.Number(value=36, label="Best Frame")
-                        out_frames = gr.Number(value=-1, label="Out Frames")
+                        best_frame = gr.Number(value=36, label="Best Frame", info="specify the frame index in the driving video where the head pose best matches the source image (note: precision of best_frame index might affect the final quality)")
+                        out_frames = gr.Number(value=-1, label="Out Frames", info="number of generation frames")
                     with gr.Accordion("Driving video Frames"):
                         driving_frames = gr.Gallery(show_label=True, columns=6, height=512, elem_id="frames-gallery")
                 with gr.Row():
@@ -143,6 +170,16 @@ with gr.Blocks(css=css) as demo:
                     ],
                     inputs=[source_image, driving_video]
                 )
+                gr.HTML("""
+                <div style="display:flex;column-gap:4px;">
+                    <a href="https://huggingface.co/spaces/fffiloni/X-Portrait?duplicate=true">
+                        <img src="https://huggingface.co/datasets/huggingface/badges/resolve/main/duplicate-this-space-xl.svg" alt="Duplicate this Space">
+                    </a>
+                    <a href="https://huggingface.co/fffiloni">
+                        <img src="https://huggingface.co/datasets/huggingface/badges/resolve/main/follow-me-on-HF-xl-dark.svg" alt="Follow me on HF">
+                    </a>
+                </div>
+                """)
 
 
     driving_video.upload(