simple_video_gen.py

"""
Simple Video Generation using SadTalker
Integrates with our Mindfull pipeline
"""
import os
import sys
import subprocess
import tempfile
import shutil
from pathlib import Path
import cv2
import numpy as np

def generate_video_simple(image_path, audio_path, output_path):
    """
    Generate video using SadTalker with fallback options
    """
    try:
        # Use SadTalker with skip_wav2lip - this should work reliably now
        print("Attempting SadTalker video generation...")
        success = generate_with_sadtalker(image_path, audio_path, output_path)
        
        if success and os.path.exists(output_path):
            size = os.path.getsize(output_path)
            print(f"✓ Video generation successful: {output_path} (size: {size} bytes)")
            return True
        else:
            print(f"✗ SadTalker failed to generate video")
            return False
        
    except Exception as e:
        print(f"✗ Video generation error: {e}")
        return False

def generate_with_sadtalker(image_path, audio_path, output_path):
    """
    Use enhanced SadTalker + Wav2Lip pipeline for video generation
    """
    try:
        # Apply patches first
        import python313_compat_patch
        
        # Try to import SadTalker components
        sadtalker_path = "sadtalker+wav2lip"
        if os.path.exists(sadtalker_path):
            # Try enhanced pipeline first
            enhanced_script = os.path.join(sadtalker_path, "enhanced_pipeline.py")
            
            if os.path.exists(enhanced_script):
                # Use enhanced pipeline with automatic quality detection
                cmd = [
                    sys.executable,
                    enhanced_script,
                    "--source_image", image_path,
                    "--audio", audio_path,
                    "--output", output_path,
                    "--use_wav2lip"  # Let the pipeline decide based on quality
                ]
            else:
                # Fallback to simple pipeline
                cmd = [
                    sys.executable,
                    os.path.join(sadtalker_path, "simple_pipeline.py"),
                    "--source_image", image_path,
                    "--audio", audio_path,
                    "--output", output_path,
                    "--skip_wav2lip"  # Skip Wav2Lip to preserve SadTalker animation
                ]

            result = subprocess.run(cmd, capture_output=True, text=True, timeout=600)
            print(f"Enhanced Pipeline STDOUT:\n{result.stdout}")
            if result.stderr:
                print(f"Enhanced Pipeline STDERR:\n{result.stderr}")

            # Check if pipeline successfully created the output file
            if result.returncode == 0 and os.path.exists(output_path):
                size = os.path.getsize(output_path)
                if size > 0:
                    print(f"✓ Enhanced video generated successfully: {output_path} (size: {size} bytes)")
                    return True
                else:
                    print(f"✗ Enhanced pipeline output file is empty: {output_path}")
                    return False
            else:
                print(f"✗ SadTalker pipeline failed or output not created")
                print(f"Return code: {result.returncode}")
                return False
        else:
            print("✗ SadTalker directory not found")
            return False
            
    except Exception as e:
        print(f"✗ SadTalker generation failed: {e}")
        return False

def generate_basic_video(image_path, audio_path, output_path):
    """
    Generate basic video by combining static image with audio
    """
    try:
        # Use FFmpeg to create video from static image + audio
        cmd = [
            "ffmpeg", "-y",
            "-loop", "1",
            "-i", image_path,
            "-i", audio_path,
            "-c:v", "libx264",
            "-c:a", "aac",
            "-b:a", "192k",
            "-shortest",
            "-pix_fmt", "yuv420p",
            output_path
        ]
        
        result = subprocess.run(cmd, capture_output=True, text=True, timeout=60)
        
        if result.returncode == 0 and os.path.exists(output_path):
            print(f"✓ Basic video generated: {output_path}")
            return True
        else:
            print(f"✗ FFmpeg failed: {result.stderr}")
            return False
            
    except subprocess.TimeoutExpired:
        print("✗ FFmpeg generation timed out")
        return False
    except Exception as e:
        print(f"✗ Basic video generation failed: {e}")
        return False

def generate_placeholder_video(image_path, audio_path, output_path):
    """
    Create a placeholder video with simple effects
    """
    try:
        import soundfile as sf
        import moviepy.editor as mp
        
        # Load audio to get duration
        try:
            audio_data, sample_rate = sf.read(audio_path)
            duration = len(audio_data) / sample_rate
        except:
            duration = 5.0  # Default 5 seconds
        
        # Load and prepare image
        image = cv2.imread(image_path)
        if image is None:
            raise ValueError(f"Could not load image: {image_path}")
        
        # Resize image to standard video dimensions
        height, width = image.shape[:2]
        target_width, target_height = 640, 480
        
        if width != target_width or height != target_height:
            image = cv2.resize(image, (target_width, target_height))
        
        # Create video with slight zoom/movement effect
        fps = 25
        total_frames = int(duration * fps)
        
        # Ensure output directory exists
        os.makedirs(os.path.dirname(output_path), exist_ok=True)
        
        # Create temporary video file
        temp_video = output_path.replace('.mp4', '_temp.mp4')
        fourcc = cv2.VideoWriter_fourcc(*'mp4v')
        out = cv2.VideoWriter(temp_video, fourcc, fps, (target_width, target_height))
        
        for frame_num in range(total_frames):
            # Add slight zoom effect
            zoom_factor = 1.0 + 0.05 * np.sin(frame_num * 0.1)
            
            if zoom_factor != 1.0:
                center_x, center_y = target_width // 2, target_height // 2
                new_width = int(target_width / zoom_factor)
                new_height = int(target_height / zoom_factor)
                
                x1 = max(0, center_x - new_width // 2)
                y1 = max(0, center_y - new_height // 2)
                x2 = min(target_width, x1 + new_width)
                y2 = min(target_height, y1 + new_height)
                
                cropped = image[y1:y2, x1:x2]
                frame = cv2.resize(cropped, (target_width, target_height))
            else:
                frame = image.copy()
            
            out.write(frame)
        
        out.release()
        
        # Combine with audio using FFmpeg
        cmd = [
            "ffmpeg", "-y",
            "-i", temp_video,
            "-i", audio_path,
            "-c:v", "libx264",
            "-c:a", "aac",
            "-shortest",
            output_path
        ]
        
        result = subprocess.run(cmd, capture_output=True, text=True)
        
        # Clean up temp video
        if os.path.exists(temp_video):
            os.remove(temp_video)
        
        if result.returncode == 0 and os.path.exists(output_path):
            print(f"✓ Placeholder video generated: {output_path}")
            return True
        else:
            print(f"✗ Placeholder video generation failed")
            return False
            
    except Exception as e:
        print(f"✗ Placeholder video generation failed: {e}")
        return False

if __name__ == "__main__":
    # Test video generation
    test_image = "avatar_assets/officer.png"
    test_audio = "outputs/test_audio_simple.wav"
    test_output = "outputs/test_video.mp4"
    
    print("Testing video generation...")
    if os.path.exists(test_image) and os.path.exists(test_audio):
        if generate_video_simple(test_image, test_audio, test_output):
            print("✓ Video generation successful")
        else:
            print("✗ Video generation failed")
    else:
        print(f"✗ Missing test files:")
        print(f"  Image: {test_image} - {'✓' if os.path.exists(test_image) else '✗'}")
        print(f"  Audio: {test_audio} - {'✓' if os.path.exists(test_audio) else '✗'}")