Update app.py

app.py

@@ -1,7 +1,8 @@
-import argparse
 import tempfile
 import os
 import shutil
+import librosa
+import json
 from flask import Flask, request, jsonify, send_from_directory
 from omegaconf import OmegaConf
 import torch
@@ -107,6 +108,43 @@ def generate_audio(voice_cloning, text_prompt):
 
     return driven_audio_path
 
+def get_video_duration(video_path):
+    """Extracts video duration dynamically using ffprobe."""
+    cmd = [
+        "ffprobe", "-v", "error", "-show_entries", "format=duration", 
+        "-of", "json", video_path
+    ]
+    result = subprocess.run(cmd, capture_output=True, text=True)
+    duration = json.loads(result.stdout)["format"]["duration"]
+    return float(duration)
+
+
+def extend_video_simple(video_path, audio_path, output_path):
+    # Get durations dynamically
+    audio_duration = librosa.get_duration(path=audio_path)
+    video_duration = get_video_duration(video_path)
+
+    print(f"Video Duration: {video_duration:.2f} sec")
+    print(f"Audio Duration: {audio_duration:.2f} sec")
+
+    if audio_duration > video_duration:
+        print("Extending video by adding reversed version.")
+
+        # Create a reversed version of the full video
+        reversed_clip = "reversed.mp4"
+        subprocess.run(
+            f"ffmpeg -i {video_path} -vf reverse -an {reversed_clip}", shell=True
+        )
+
+        # Merge original + reversed
+        subprocess.run(
+            f"ffmpeg -i {video_path} -i {reversed_clip} -filter_complex \"[0:v:0][1:v:0]concat=n=2:v=1[outv]\" -map \"[outv]\" -an {output_path}",
+            shell=True
+        )
+
+    else:
+        print("Audio is not longer than video. No extension needed.")
+        subprocess.run(f"cp {video_path} {output_path}", shell=True)
 
 @app.route('/run', methods=['POST'])
 def generate_video():
@@ -136,6 +174,8 @@ def generate_video():
     # You can pass additional parameters via form data if needed (e.g., checkpoint path)
     inference_ckpt_path = request.form.get('inference_ckpt_path', 'checkpoints/latentsync_unet.pt')
     unet_config_path = request.form.get('unet_config_path', 'configs/unet/second_stage.yaml')
+
+    extend_video_simple(temp_video_path, temp_audio_path, output_video)
    
     try:
         run_inference(