Update app.py

app.py

@@ -55,11 +55,35 @@ OUTPUT_DIR = "gradio_output"
 WEIGHTS_DIR = "pretrain_weights"
 REPO_ID = "lixinyizju/moda"
 
+# --- Automatic Config Patching for Aspect Ratio Preservation ---
+def patch_liveportrait_config():
+    """
+    Patches the liveportrait_config.yaml to enable 'flag_pasteback'.
+    This stitches the animated face back onto the original high-resolution full-sized
+    source image, maintaining the exact original image size and aspect ratio.
+    """
+    config_path = "configs/audio2motion/model/liveportrait_config.yaml"
+    if os.path.exists(config_path):
+        try:
+            import yaml
+            with open(config_path, 'r', encoding='utf-8') as f:
+                config_data = yaml.safe_load(f) or {}
+            
+            # Enable pasteback to restore original image dimensions in the output video
+            if config_data.get('flag_pasteback') is not True:
+                config_data['flag_pasteback'] = True
+                with open(config_path, 'w', encoding='utf-8') as f:
+                    yaml.dump(config_data, f)
+                print("Successfully patched liveportrait_config.yaml: flag_pasteback set to True")
+        except Exception as e:
+            print(f"Warning: Failed to automatically patch configuration file: {e}")
+
 # --- Download Pre-trained Weights from Hugging Face Hub ---
 def download_weights():
     """
     Downloads pre-trained weights from Hugging Face Hub if they don't exist locally.
     """
+    # A simple check for a key file to see if the download is likely complete
     motion_model_file = os.path.join(WEIGHTS_DIR, "moda", "net-200.pth")
     
     if not os.path.exists(motion_model_file):
@@ -96,16 +120,20 @@ def ensure_wav_format(audio_path, max_duration_ms=300000): # Increased to 5 minu
     
     try:
         print(f"Loading and processing audio from: {audio_path}")
+        # Load the audio file (supports wav, mp3, m4a, etc.)
         audio = AudioSegment.from_file(audio_path)
         
+        # Trim if the audio exceeds the limit
         if len(audio) > max_duration_ms:
             print(f"Audio is {len(audio)/1000:.2f}s, trimming to {max_duration_ms/1000:.2f}s")
             audio = audio[:max_duration_ms]
         else:
             print(f"Audio is {len(audio)/1000:.2f}s (under the {max_duration_ms/1000:.2f}s limit)")
         
+        # Create a temporary WAV file
         with tempfile.NamedTemporaryFile(suffix='.wav', delete=False) as tmp_file:
             wav_path = tmp_file.name
+            # Export as WAV with standard settings
             audio.export(
                 wav_path,
                 format='wav',
@@ -147,57 +175,62 @@ def extract_last_frame(video_path, output_image_path):
 
 def concatenate_videos_ffmpeg(video_paths, full_audio_path, output_path):
     """
-    Stitches multiple video files together using FFMPEG concat filter (with full re-encoding).
-    Standardizes dimensions to 512x512, sets SAR to 1, and forces a strict 25 FPS timebase.
-    Finally overlays the clean original audio track.
+    Stitches multiple video files together (video track only) and overlays the original full audio track.
+    This bypasses any pops or sync problems at the chunk boundaries.
     """
+    # Create the text file listing all videos
+    with tempfile.NamedTemporaryFile(mode='w', suffix='.txt', delete=False, encoding='utf-8') as f:
+        for p in video_paths:
+            # Format absolute path with forward slashes and escape single quotes
+            abs_path = os.path.abspath(p).replace('\\', '/')
+            safe_path = abs_path.replace("'", "'\\''")
+            f.write(f"file '{safe_path}'\n")
+        list_file_path = f.name
+        
     temp_video_only = os.path.join(os.path.dirname(output_path), "temp_video_only.mp4")
-    num_inputs = len(video_paths)
         
     try:
-        # Build the FFMPEG command with input files
-        cmd_concat = ['ffmpeg', '-y']
-        for p in video_paths:
-            cmd_concat.extend(['-i', str(p)])
-            
-        # Construct the complex filter: scale to 512x512, set SAR, and concatenate
-        filter_parts = []
-        concat_inputs = []
-        for i in range(num_inputs):
-            filter_parts.append(f"[{i}:v]scale=512:512,setsar=1[v{i}]")
-            concat_inputs.append(f"[v{i}]")
-            
-        concat_inputs_str = "".join(concat_inputs)
-        filter_parts.append(f"{concat_inputs_str}concat=n={num_inputs}:v=1:a=0[v]")
-        filter_complex_str = "; ".join(filter_parts)
-        
-        # Output setup for high compatibility (H.264, YUV420p, strict 25 FPS)
-        cmd_concat.extend([
-            '-filter_complex', filter_complex_str,
-            '-map', '[v]',
-            '-c:v', 'libx264',
-            '-preset', 'medium',
-            '-crf', '18',           # High quality output
-            '-pix_fmt', 'yuv420p',  # High compatibility pixel format
-            '-r', '25',             # Strictly 25 fps matching default output
+        # Step 1: Concatenate video tracks only (discard original segment audios to avoid jumps)
+        cmd_copy = [
+            'ffmpeg', '-y',
+            '-f', 'concat',
+            '-safe', '0',
+            '-i', list_file_path,
+            '-an',             # Discard segment audios completely
+            '-c:v', 'copy',    # Direct stream copy (instant & lossless)
             temp_video_only
-        ])
+        ]
+        print(f"Executing lossless FFMPEG video concat (no-audio): {' '.join(cmd_copy)}")
+        result = subprocess.run(cmd_copy, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
         
-        print(f"Executing video concat filter: {' '.join(cmd_concat)}")
-        result_concat = subprocess.run(cmd_concat, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
-        if result_concat.returncode != 0:
-            raise Exception(f"FFMPEG video-only concat filter failed: {result_concat.stderr.decode('utf-8')}")
+        if result.returncode != 0:
+            print("Lossless video-only concatenation failed. Retrying with full video re-encoding...")
+            # Fallback path: video re-encoding (robust)
+            cmd_reencode = [
+                'ffmpeg', '-y',
+                '-f', 'concat',
+                '-safe', '0',
+                '-i', list_file_path,
+                '-an',
+                '-c:v', 'libx264',
+                '-pix_fmt', 'yuv420p',
+                temp_video_only
+            ]
+            print(f"Executing FFMPEG video-only re-encode: {' '.join(cmd_reencode)}")
+            result_re = subprocess.run(cmd_reencode, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
+            if result_re.returncode != 0:
+                raise Exception(f"FFMPEG video concatenation failed: {result_re.stderr.decode('utf-8')}")
         
-        # Step 2: Merge the newly rendered video with the complete original audio file
+        # Step 2: Merge the concatenated video track with the complete original audio file
         cmd_merge = [
             'ffmpeg', '-y',
             '-i', temp_video_only,
             '-i', str(full_audio_path),
-            '-map', '0:v:0',    # Map video from temp file
-            '-map', '1:a:0',    # Map audio from original full wav file
-            '-c:v', 'copy',     # Direct copy of the already encoded video
+            '-map', '0:v:0',    # Use video from concatenated file
+            '-map', '1:a:0',    # Use audio from original full wav file
+            '-c:v', 'copy',     # Keep video intact
             '-c:a', 'aac',      # Encode original audio to standard clean AAC
-            '-shortest',        # Sync to shortest stream
+            '-shortest',        # Sync to the shortest stream duration
             str(output_path)
         ]
         print(f"Applying final original audio track: {' '.join(cmd_merge)}")
@@ -206,7 +239,9 @@ def concatenate_videos_ffmpeg(video_paths, full_audio_path, output_path):
             raise Exception(f"FFMPEG audio replacement failed: {result_merge.stderr.decode('utf-8')}")
                 
     finally:
-        # Clean up temporary video file
+        # Clean up temporary files
+        if os.path.exists(list_file_path):
+            os.remove(list_file_path)
         if os.path.exists(temp_video_only):
             try:
                 os.remove(temp_video_only)
@@ -217,6 +252,9 @@ def concatenate_videos_ffmpeg(video_paths, full_audio_path, output_path):
 # Create output directory if it doesn't exist
 os.makedirs(OUTPUT_DIR, exist_ok=True)
 
+# Patch the LivePortrait config to enable full-frame pasteback
+patch_liveportrait_config()
+
 # Download weights before initializing the pipeline
 download_weights()
 
@@ -350,12 +388,12 @@ def generate_motion(source_image_path, driving_audio_path, emotion_name, cfg_sca
                     print("Falling back to the original source image for the next segment.")
                     current_image_path = source_image_path
 
-        # 4. Merge all generated chunks and overlay the original full audio track
+        # 4. Merge all generated chunks and lay the original full audio track on top
         progress(0.95, desc="Merging segments and applying original audio...")
         final_video_name = f"final_{timestamp}.mp4"
         final_path = Path(os.path.join(run_output_dir, final_video_name))
         
-        # Force re-encoding merge to ensure absolute smooth transitions without lag
+        # We always run this function (even for 1 chunk) to ensure the original, clean, high-quality WAV audio is perfectly mapped on the video
         concatenate_videos_ffmpeg(video_chunks_paths, wav_audio_path, final_path)
             
     except Exception as e: