2024-10-07 10:18:37,933 [ERROR] Mouth frame generation returned no frames.import time
import os
import uuid
import cv2
import tqdm
import shutil
import sys
import re
import wave
from piper.voice import PiperVoice
from talkingface.audio_model import AudioModel
from talkingface.render_model import RenderModel
import subprocess
from concurrent.futures import ProcessPoolExecutor, as_completed
import multiprocessing
import warnings
import logging

# =============================================================
#                       Setup Logging
# =============================================================
logging.basicConfig(
    filename='processing.log',
    filemode='a',
    format='%(asctime)s [%(levelname)s] %(message)s',
    level=logging.DEBUG  # Changed to DEBUG for more detailed logs
)

# Suppress specific warnings (use with caution)
warnings.filterwarnings("ignore", category=UserWarning, module="sklearn")

# =============================================================
#                      Helper Functions
# =============================================================

def text_preprocess(text):
    """
    Preprocess the input text by removing specific patterns and non-ASCII characters.
    """
    logging.info("Text before preprocessing: %s", text)
    text = re.sub(r'~+', '!', text)
    text = re.sub(r"\(.*?\)", "", text)
    text = re.sub(r"(\*[^*]+\*)|(_[^_]+_)", "", text).strip()
    text = re.sub(r'[^\x00-\x7F]+', '', text)
    logging.info("Preprocessed Text: %s", text)
    return text

def run_ffmpeg_command(cmd, task_id):
    """
    Execute an FFmpeg command using subprocess.run and handle errors.
    """
    try:
        logging.info(f"task_id: {task_id}. Running FFmpeg command: {cmd}")
        result = subprocess.run(cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
        if result.returncode != 0:
            logging.error(f"task_id: {task_id}. FFmpeg error: {result.stderr}")
            return False
        logging.debug(f"task_id: {task_id}. FFmpeg command executed successfully.")
        return True
    except Exception as e:
        logging.error(f"task_id: {task_id}. FFmpeg execution failed: {str(e)}")
        return False

def synthesize_audio(text, output_wav_path):
    """
    Synthesize speech from text and save it to a WAV file.
    """
    model = "en_US-john-medium.onnx"
    try:
        voice = PiperVoice.load(model)
        logging.info("PiperVoice model loaded successfully.")
    except Exception as e:
        logging.error(f"Failed to load PiperVoice model: {str(e)}")
        sys.exit(1)
    
    try:
        with wave.open(output_wav_path, 'w') as wav_file:
            voice.synthesize(text_preprocess(text), wav_file)
        logging.info(f"Audio synthesized and saved to {output_wav_path}")
    except Exception as e:
        logging.error(f"Audio synthesis failed: {str(e)}")
        sys.exit(1)
    return output_wav_path

def convert_audio(input_wav, output_wav, task_id):
    """
    Convert audio to the desired format using FFmpeg.
    """
    ffmpeg_cmd = f"ffmpeg -y -i \"{input_wav}\" -ar 16000 \"{output_wav}\""
    success = run_ffmpeg_command(ffmpeg_cmd, task_id)
    if success:
        logging.info(f"Audio converted and saved to {output_wav}")
    else:
        logging.error(f"Audio conversion failed for {input_wav}")
    return success

def process_frame(frame, render_model_path):
    """
    Process a single video frame using the RenderModel.
    Each process initializes its own RenderModel instance.
    """
    try:
        from talkingface.render_model import RenderModel  # Ensure import within worker
        renderModel = RenderModel()
        renderModel.loadModel(render_model_path)
        logging.debug("RenderModel loaded successfully in worker.")
        
        if frame is None:
            logging.error("Received None frame for processing.")
            return None
        
        processed_frame = renderModel.interface(frame)
        if processed_frame is None:
            logging.error("RenderModel.interface returned None for a frame.")
            return None
        
        return processed_frame
    except AttributeError as ae:
        logging.error(f"AttributeError in process_frame: {str(ae)}")
        return None
    except Exception as e:
        logging.error(f"Unexpected error in process_frame: {str(e)}")
        return None

def merge_audio_video(save_path, wavpath, output_video_name, task_id):
    """
    Merge processed video and audio into the final output video using FFmpeg.
    """
    ffmpeg_cmd = (
        f"ffmpeg -y -i \"{save_path}\" -i \"{wavpath}\" "
        f"-c:v libx264 -pix_fmt yuv420p -c:a aac -strict experimental "
        f"-loglevel error \"{output_video_name}\""
    )
    success = run_ffmpeg_command(ffmpeg_cmd, task_id)
    if success:
        logging.info(f"Audio and video merged into {output_video_name}")
    else:
        logging.error(f"Merging audio and video failed for {output_video_name}")
    return success

# =============================================================
#                          Main Function
# =============================================================

def main():
    # Initialize AudioModel
    video_input_path = sys.argv[1]
    output_video_name = sys.argv[3]
    audioModel = AudioModel()
    try:
        audioModel.loadModel("checkpoint/audio.pkl")
        logging.info("AudioModel loaded successfully.")
    except Exception as e:
        logging.error(f"Failed to load AudioModel: {str(e)}")
        sys.exit(1)

    # Define paths and filenames
    filename = 'temp.wav'
    text = '''Gunman lurked for hours before Trump's last-minute game of golf
I think he’ll still make a presence, not just here but everywhere. I don't think that’s something he’s going to be shying away from.”'''
    
    # Synthesize audio
    wav_file = synthesize_audio(text, filename)

    # Ensure correct number of command-line arguments
   

   

    logging.info("=====================")
    logging.info(f"Video input path: {video_input_path}")
    logging.info(f"Output video name: {output_video_name}")

    audio_path = filename
    wavpath = 'output_audio1.wav'

    task_id = str(uuid.uuid1())
    logging.info(f"Generated task_id: {task_id}")

    # Convert audio
    if not convert_audio(audio_path, wavpath, task_id):
        logging.error(f"task_id: {task_id}. Audio conversion failed.")
        sys.exit(1)

    # Prepare video with RenderModel (main process)
    try:
        renderModel_main = RenderModel()
        renderModel_main.loadModel("checkpoint/render.pth")
        logging.info("RenderModel loaded successfully in main process.")
        
        pkl_path = os.path.join(video_input_path, "keypoint_rotate.pkl")
        video_path_processed = os.path.join(video_input_path, "circle.mp4")

        if not os.path.exists(pkl_path):
            logging.error(f"Keypoint pickle file not found at {pkl_path}")
            sys.exit(1)
        
        if not os.path.exists(video_path_processed):
            logging.error(f"Processed video file not found at {video_path_processed}")
            sys.exit(1)
        
        renderModel_main.reset_charactor(video_path_processed, pkl_path)
        logging.info("RenderModel.reset_charactor completed successfully.")
    except Exception as e:
        logging.error(f"Failed to reset_charactor: {str(e)}")
        sys.exit(1)

    # Generate mouth frames from audio
    try:
        if not os.path.exists(wavpath):
            logging.error(f"Mouth frame generation failed: audio file {wavpath} does not exist.")
            sys.exit(1)
        
        # Check if audio file has non-zero length
        if os.path.getsize(wavpath) == 0:
            logging.error(f"Mouth frame generation failed: audio file {wavpath} is empty.")
            sys.exit(1)
        
        # Log properties of the audio file
        with wave.open(wavpath, 'rb') as wav_file_obj:
            n_channels = wav_file_obj.getnchannels()
            sampwidth = wav_file_obj.getsampwidth()
            framerate = wav_file_obj.getframerate()
            n_frames = wav_file_obj.getnframes()
            params = wav_file_obj.getparams()
            logging.info(f"Audio file properties - Channels: {n_channels}, Sample Width: {sampwidth}, "
                         f"Frame Rate: {framerate} Hz, Number of Frames: {n_frames}")
        
        mouth_frame = audioModel.interface_wav(wavpath)
        if not isinstance(mouth_frame, list) or not mouth_frame:
            logging.error("Mouth frame generation returned no frames.")
            sys.exit(1)
        logging.info(f"Generated {len(mouth_frame)} mouth frames from audio.")
    except Exception as e:
        logging.error(f"Failed to generate mouth frames: {str(e)}")
        sys.exit(1)

    # Read video properties
    cap_input = cv2.VideoCapture(video_path_processed)
    if not cap_input.isOpened():
        logging.error(f"task_id: {task_id}. Unable to open video: {video_path_processed}")
        sys.exit(1)
    vid_width = int(cap_input.get(cv2.CAP_PROP_FRAME_WIDTH))
    vid_height = int(cap_input.get(cv2.CAP_PROP_FRAME_HEIGHT))
    cap_input.release()
    logging.info(f"Video properties - Width: {vid_width}, Height: {vid_height}")

    # Prepare output directory and video writer
    output_dir = os.path.join("output", task_id)
    os.makedirs(output_dir, exist_ok=True)
    save_path = os.path.join(output_dir, "silence.mp4")

    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    videoWriter = cv2.VideoWriter(save_path, fourcc, 25, (vid_width, vid_height))
    if not videoWriter.isOpened():
        logging.error(f"Failed to initialize VideoWriter with path: {save_path}")
        shutil.rmtree(output_dir)
        sys.exit(1)
    logging.info(f"VideoWriter initialized with save path: {save_path}")

    # Process frames with concurrency
    logging.info("Starting frame processing with concurrency...")
    start_time = time.time()

    render_model_path = "checkpoint/render.pth"

    try:
        max_workers = multiprocessing.cpu_count() or 4
        logging.info(f"Using {max_workers} worker processes for frame processing.")
        
        with ProcessPoolExecutor(max_workers=max_workers) as executor:
            # Submit all frame processing tasks
            future_to_frame = {
                executor.submit(process_frame, frame, render_model_path): idx 
                for idx, frame in enumerate(mouth_frame)
            }

            for future in tqdm.tqdm(as_completed(future_to_frame), total=len(future_to_frame), desc="Processing Frames"):
                idx = future_to_frame[future]
                try:
                    processed_frame = future.result()
                    if processed_frame is not None:
                        videoWriter.write(processed_frame)
                        logging.debug(f"Frame {idx} processed and written successfully.")
                    else:
                        logging.warning(f"task_id: {task_id}. Received None for processed frame {idx}. Skipping frame.")
                except Exception as e:
                    logging.error(f"task_id: {task_id}. Frame processing failed for frame {idx}: {str(e)}")
                    videoWriter.release()
                    shutil.rmtree(output_dir)
                    sys.exit(1)
    except Exception as e:
        logging.error(f"task_id: {task_id}. An error occurred during frame processing: {str(e)}")
        videoWriter.release()
        shutil.rmtree(output_dir)
        sys.exit(1)

    videoWriter.release()
    elapsed_time = time.time() - start_time
    logging.info(f"Frame processing completed in {elapsed_time:.2f} seconds.")

    # Merge audio and video
    final_output_path = os.path.abspath(output_video_name)
    if not merge_audio_video(save_path, wavpath, final_output_path, task_id):
        logging.error(f"task_id: {task_id}. Merging audio and video failed.")
        shutil.rmtree(output_dir)
        sys.exit(1)

    logging.info(f"Output video saved to {final_output_path}")

    # Clean up temporary files and directories
    try:
        shutil.rmtree(output_dir)
        logging.info(f"Cleaned up temporary directory: {output_dir}")
    except Exception as e:
        logging.warning(f"Failed to remove temporary directory {output_dir}: {str(e)}")

    try:
        if os.path.exists(wavpath):
            os.remove(wavpath)
            logging.info(f"Removed temporary audio file: {wavpath}")
        if os.path.exists(audio_path):
            os.remove(audio_path)
            logging.info(f"Removed temporary audio file: {audio_path}")
    except Exception as e:
        logging.warning(f"Failed to remove temporary files: {str(e)}")

# =============================================================
#                        Entry Point
# =============================================================

if __name__ == "__main__":
    # Set the multiprocessing start method to 'spawn' to avoid CUDA errors
    try:
        multiprocessing.set_start_method('spawn', force=True)
        logging.info("Multiprocessing start method set to 'spawn'.")
    except RuntimeError as e:
        # Start method has already been set
        logging.warning(f"Could not set start method to 'spawn': {str(e)}")
    main()