nota_wav2lip/demo.py

import argparse
import platform
import subprocess
import time
from pathlib import Path
from typing import Dict, Iterator, List, Literal, Optional, Union

import cv2
import numpy as np

from config import hparams as hp
from nota_wav2lip.inference import Wav2LipInferenceImpl
from nota_wav2lip.util import FFMPEG_LOGGING_MODE
from nota_wav2lip.video import AudioSlicer, VideoSlicer


class Wav2LipModelComparisonDemo:
    def __init__(self, device='cpu', result_dir='./temp', model_list: Optional[Union[str, List[str]]]=None):
        if model_list is None:
            model_list: List[str] = ['wav2lip', 'nota_wav2lip']
        if isinstance(model_list, str) and len(model_list) != 0:
            model_list: List[str] = [model_list]
        super().__init__()
        self.video_dict: Dict[str, VideoSlicer] = {}
        self.audio_dict: Dict[str, AudioSlicer] = {}

        self.model_zoo: Dict[str, Wav2LipInferenceImpl] = {}
        for model_name in model_list:
            assert model_name in hp.inference.model, f"{model_name} not in hp.inference_model: {hp.inference.model}"
            self.model_zoo[model_name] = Wav2LipInferenceImpl(
                model_name, hp_inference_model=hp.inference.model[model_name], device=device
            )

        self._params_zoo: Dict[str, str] = {
            model_name: self.model_zoo[model_name].params for model_name in self.model_zoo
        }

        self.result_dir: Path = Path(result_dir)
        self.result_dir.mkdir(exist_ok=True)

    @property
    def params(self):
        return self._params_zoo

    def _infer(
        self,
        audio_name: str,
        video_name: str,
        model_type: Literal['wav2lip', 'nota_wav2lip']
    ) -> Iterator[np.ndarray]:
        audio_iterable: AudioSlicer = self.audio_dict[audio_name]
        video_iterable: VideoSlicer = self.video_dict[video_name]
        target_model = self.model_zoo[model_type]
        return target_model.inference_with_iterator(audio_iterable, video_iterable)

    def update_audio(self, audio_path, name=None):
        _name = name if name is not None else Path(audio_path).stem
        self.audio_dict.update(
            {_name: AudioSlicer(audio_path)}
        )

    def update_video(self, frame_dir_path, bbox_path, name=None):
        _name = name if name is not None else Path(frame_dir_path).stem
        self.video_dict.update(
            {_name: VideoSlicer(frame_dir_path, bbox_path)}
        )

    def save_as_video(self, audio_name, video_name, model_type):

        output_video_path = self.result_dir / 'generated_with_audio.mp4'
        frame_only_video_path = self.result_dir / 'generated.mp4'
        audio_path = self.audio_dict[audio_name].audio_path

        out = cv2.VideoWriter(str(frame_only_video_path),
                              cv2.VideoWriter_fourcc(*'mp4v'),
                              hp.face.video_fps,
                              (hp.inference.frame.w, hp.inference.frame.h))
        start = time.time()
        for frame in self._infer(audio_name=audio_name, video_name=video_name, model_type=model_type):
            out.write(frame)
        inference_time = time.time() - start
        out.release()

        command = f"ffmpeg {FFMPEG_LOGGING_MODE['ERROR']} -y -i {audio_path} -i {frame_only_video_path} -strict -2 -q:v 1 {output_video_path}"
        subprocess.call(command, shell=platform.system() != 'Windows')

        # The number of frames of generated video
        video_frames_num = len(self.audio_dict[audio_name])
        inference_fps = video_frames_num / inference_time

        return output_video_path, inference_time, inference_fps