ANI_plus.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "source": [
        "#%cd /content/\n",
        "!git clone https://github.com/G-force78/AniPortrait-plus.git\n"
      ],
      "metadata": {
        "id": "oQnF-glg0IFd"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title install reqs\n",
        "!apt install espeak-ng aria2\n",
        "!aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/camenduru/AniPortrait/resolve/main/audio2mesh.pt -o /Aniportrait-plus/pretrained_model/audio2mesh.pt\n",
        "!aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/camenduru/AniPortrait/resolve/main/denoising_unet.pth -o /Aniportrait-plus/pretrained_model/denoising_unet.pth\n",
        "!aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/camenduru/AniPortrait/resolve/main/motion_module.pth -o /Aniportrait-plus/pretrained_model/motion_module.pth\n",
        "!aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/camenduru/AniPortrait/resolve/main/pose_guider.pth -o /Aniportrait-plus/pretrained_model/pose_guider.pth\n",
        "!aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/camenduru/champ/resolve/main/champ/denoising_unet.pth -o /Aniportrait-plus/pretrained_model/denoising_unet.pth\n",
        "!aria2c  --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/ZJYang/AniPortrait/resolve/main/film_net_fp16.pt -o /Aniportrait-plus/pretrained_model/film_net_fp16.pt\n",
        "!aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/camenduru/champ/resolve/main/champ/guidance_encoder_depth.pth -o /Aniportrait-plus/pretrained_model/guidance_encoder_depth.pth\n",
        "!aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/camenduru/champ/resolve/main/champ/guidance_encoder_dwpose.pth -o /Aniportrait-plus/pretrained_model/guidance_encoder_dwpose.pth\n",
        "!aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/camenduru/champ/resolve/main/champ/guidance_encoder_normal.pth -o /Aniportrait-plus/pretrained_model/guidance_encoder_normal.pth\n",
        "!aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/camenduru/champ/resolve/main/champ/guidance_encoder_semantic_map.pth -o /Aniportrait-plus/pretrained_model//guidance_encoder_semantic_map.pth\n",
        "!aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/camenduru/champ/resolve/main/champ/motion_module.pth -o /Aniportrait-plus/pretrained_model/motion_module.pth\n",
        "!aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/camenduru/champ/resolve/main/champ/reference_unet.pth -o /Aniportrait-plus/pretrained_model/reference_unet.pth\n",
        "!aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/camenduru/champ/raw/main/image_encoder/config.json -o /Aniportrait-plus/pretrained_model/image_encoder/config.json\n",
        "!aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/camenduru/champ/resolve/main/image_encoder/pytorch_model.bin -o /Aniportrait-plus/pretrained_model/image_encoder/pytorch_model.bin\n",
        "!aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/camenduru/champ/raw/main/sd-vae-ft-mse/config.json -o /Aniportrait-plus/pretrained_model/sd-vae-ft-mse/config.json\n",
        "!aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/camenduru/champ/resolve/main/sd-vae-ft-mse/diffusion_pytorch_model.bin -o /Aniportrait-plus/pretrained_model/sd-vae-ft-mse/diffusion_pytorch_model.bin\n",
        "!aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/camenduru/champ/raw/main/stable-diffusion-v1-5/feature_extractor/preprocessor_config.json -o /Aniportrait-plus/pretrained_model/stable-diffusion-v1-5/feature_extractor/preprocessor_config.json\n",
        "!aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/camenduru/champ/raw/main/stable-diffusion-v1-5/model_index.json -o /Aniportrait-plus/pretrained_model/stable-diffusion-v1-5/model_index.json\n",
        "!aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/camenduru/champ/raw/main/stable-diffusion-v1-5/unet/config.json -o /Aniportrait-plus/pretrained_model/stable-diffusion-v1-5/unet/config.json\n",
        "!aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/camenduru/champ/resolve/main/stable-diffusion-v1-5/unet/diffusion_pytorch_model.bin -o /Aniportrait-plus/pretrained_model/stable-diffusion-v1-5/unet/diffusion_pytorch_model.bin\n",
        "!aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/camenduru/champ/raw/main/stable-diffusion-v1-5/v1-inference.yaml -o /Aniportrait-plus/pretrained_model/stable-diffusion-v1-5/v1-inference.yaml\n",
        "!aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/facebook/wav2vec2-base-960h/raw/main/config.json -o /Aniportrait-plus/pretrained_model/wav2vec2-base-960h/config.json\n",
        "!aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/facebook/wav2vec2-base-960h/raw/main/feature_extractor_config.json -o /Aniportrait-plus/pretrained_model/wav2vec2-base-960h/feature_extractor_config.json\n",
        "!aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/facebook/wav2vec2-base-960h/raw/main/preprocessor_config.json -o /Aniportrait-plus/pretrained_model/wav2vec2-base-960h/preprocessor_config.json\n",
        "!aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/facebook/wav2vec2-base-960h/resolve/main/pytorch_model.bin -o /Aniportrait-plus/pretrained_model/wav2vec2-base-960h/pytorch_model.bin\n",
        "!aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/facebook/wav2vec2-base-960h/raw/main/special_tokens_map.json -o /Aniportrait-plus/pretrained_model/wav2vec2-base-960h/special_tokens_map.json\n",
        "!aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/facebook/wav2vec2-base-960h/resolve/main/tf_model.h5 -o /Aniportrait-plus/pretrained_model/wav2vec2-base-960h//tf_model.h5\n",
        "!aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/facebook/wav2vec2-base-960h/raw/main/tokenizer_config.json -o /Aniportrait-plus/pretrained_model/wav2vec2-base-960h/tokenizer_config.json\n",
        "!aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/facebook/wav2vec2-base-960h/raw/main/vocab.json -o /Aniportrait-plus/pretrained_model/wav2vec2-base-960h/vocab.json\n",
        "\n",
        "!pip install -q imageio-ffmpeg==0.4.9 onnx gradio onnxruntime GFPGAN ffmpeg-python==0.2.0 av==11.0.0 omegaconf==2.2.3 diffusers==0.24.0 mediapipe==0.10.11 einops==0.4.1 accelerate==0.21.0 xformers==0.0.25 librosa==0.9.2\n",
        "!pip install onnxruntime-gpu==1.17 realESRGAN"
      ],
      "metadata": {
        "id": "OobsWv6LtcOW",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title download upscale models\n",
        "\n",
        "\n",
        "!wget https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.4.pth -O /content/AniPortrait-plus/pretrained_model/GFPGANv1.4.pth\n",
        "\n",
        "\n",
        "\n",
        "!wget https://huggingface.co/ai-forever/Real-ESRGAN/resolve/main/RealESRGAN_x2.pth -O /content/AniPortrait-plus/pretrained_model/RealESRGAN_x2.pth\n",
        "\n",
        "!wget https://huggingface.co/ai-forever/Real-ESRGAN/resolve/main/RealESRGAN_x4.pth -O /content/AniPortrait-plus/pretrained_model/RealESRGAN_x4.pth\n",
        "\n",
        "!wget https://huggingface.co/ai-forever/Real-ESRGAN/resolve/main/RealESRGAN_x8.pth -O /content/AniPortrait-plus/pretrained_model/RealESRGAN_x8.pth\n",
        "\n",
        "!wget https://huggingface.co/bluefoxcreation/Codeformer-ONNX/resolve/main/codeformer.onnx -O /content/AniPortrait-plus/pretrained_model/codeformer.onnx\n"
      ],
      "metadata": {
        "id": "iPjXr-JytQ32",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title do this before launching app\n",
        "\n",
        "Open /usr/local/lib/python3.10/dist-packages/basicsr/data/degradations.py and on line 8, simply change:\n",
        "\n",
        "from torchvision.transforms.functional_tensor import rgb_to_grayscale\n",
        "to:\n",
        "\n",
        "from torchvision.transforms.functional import rgb_to_grayscale"
      ],
      "metadata": {
        "id": "LyvHFwEXxQYk",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Set the desired temporary directory path\n",
        "temp_dir = \"/content/AniPortrait-plus/gradiotmp\"\n",
        "os.environ[\"GRADIO_TEMP_DIR\"] = temp_dir\n",
        "%cd /content/AniPortrait-plus\n",
        "!python /content/AniPortrait-plus/aniplus.py"
      ],
      "metadata": {
        "id": "JCaFXDd-I-Hp"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "!python aniplus.py"
      ],
      "metadata": {
        "id": "ErtqIN8RhWce"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title FINAL WORKING FAST APP WITH UPSCALER\n",
        "import gradio as gr\n",
        "import os\n",
        "import shutil\n",
        "import ffmpeg\n",
        "from datetime import datetime\n",
        "from pathlib import Path\n",
        "import numpy as np\n",
        "import cv2\n",
        "import torch\n",
        "import onnxruntime\n",
        "import onnxruntime as ort\n",
        "from diffusers import AutoencoderKL, DDIMScheduler\n",
        "from einops import repeat\n",
        "from omegaconf import OmegaConf\n",
        "from PIL import Image\n",
        "from torchvision import transforms\n",
        "from transformers import CLIPVisionModelWithProjection\n",
        "from scipy.interpolate import interp1d\n",
        "from tqdm import tqdm\n",
        "from face_enhancer import (\n",
        "    get_available_enhancer_names,\n",
        "    load_face_enhancer_model,\n",
        "    cv2_interpolations,\n",
        ")\n",
        "import tempfile\n",
        "import uuid\n",
        "import os\n",
        "import cv2\n",
        "import ffmpeg\n",
        "from face_enhancer import load_face_enhancer_model\n",
        "\n",
        "from tqdm import tqdm\n",
        "\n",
        "\n",
        "from src.models.pose_guider import PoseGuider\n",
        "from src.models.unet_2d_condition import UNet2DConditionModel\n",
        "from src.models.unet_3d import UNet3DConditionModel\n",
        "from src.pipelines.pipeline_pose2vid_long import Pose2VideoPipeline\n",
        "from src.utils.util import get_fps, read_frames, save_videos_grid\n",
        "\n",
        "from src.audio_models.model import Audio2MeshModel\n",
        "from src.utils.audio_util import prepare_audio_feature\n",
        "from src.utils.mp_utils  import LMKExtractor\n",
        "from src.utils.draw_util import FaceMeshVisualizer\n",
        "from src.utils.pose_util import project_points, project_points_with_trans, matrix_to_euler_and_translation, euler_and_translation_to_matrix\n",
        "from src.utils.util import crop_face\n",
        "from scripts.vid2vid import smooth_pose_seq\n",
        "from src.utils.frame_interpolation import init_frame_interpolation_model, batch_images_interpolation_tool\n",
        "\n",
        "\n",
        "## ------------------------------ FACE ENHANCEMENT ------------------------------\n",
        "#to be continued\n",
        "PROVIDER = [\"CUDAExecutionProvider\"]\n",
        "\n",
        "\n",
        "available_providers = onnxruntime.get_available_providers()\n",
        "if \"CUDAExecutionProvider\" in available_providers:\n",
        "  print(\"\\n********** Running on CUDA **********\\n\")\n",
        "#PROVIDER = [\"CUDAExecutionProvider\", \"CPUExecutionProvider\"]\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "device = \"cuda\"\n",
        "EMPTY_CACHE = lambda: torch.cuda.empty_cache() if device == \"cuda\" else None\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "config = OmegaConf.load('/content/AniPortrait-plus/configs/prompts/animation_audio.yaml')\n",
        "if config.weight_dtype == \"fp16\":\n",
        "    weight_dtype = torch.float16\n",
        "else:\n",
        "    weight_dtype = torch.float32\n",
        "\n",
        "audio_infer_config = OmegaConf.load(config.audio_inference_config)\n",
        "# prepare model\n",
        "a2m_model = Audio2MeshModel(audio_infer_config['a2m_model'])\n",
        "a2m_model.load_state_dict(torch.load(audio_infer_config['pretrained_model']['a2m_ckpt'], map_location=\"cpu\"), strict=False)\n",
        "a2m_model.cuda().eval()\n",
        "\n",
        "vae = AutoencoderKL.from_pretrained(\n",
        "    config.pretrained_vae_path,\n",
        ").to(\"cuda\", dtype=weight_dtype)\n",
        "\n",
        "reference_unet = UNet2DConditionModel.from_pretrained(\n",
        "    config.pretrained_base_model_path,\n",
        "    subfolder=\"unet\",\n",
        ").to(dtype=weight_dtype, device=\"cuda\")\n",
        "\n",
        "inference_config_path = config.inference_config\n",
        "infer_config = OmegaConf.load(inference_config_path)\n",
        "denoising_unet = UNet3DConditionModel.from_pretrained_2d(\n",
        "    config.pretrained_base_model_path,\n",
        "    config.motion_module_path,\n",
        "    subfolder=\"unet\",\n",
        "    unet_additional_kwargs=infer_config.unet_additional_kwargs,\n",
        ").to(dtype=weight_dtype, device=\"cuda\")\n",
        "\n",
        "pose_guider = PoseGuider(noise_latent_channels=320, use_ca=True).to(device=\"cuda\", dtype=weight_dtype) # not use cross attention\n",
        "\n",
        "image_enc = CLIPVisionModelWithProjection.from_pretrained(\n",
        "    config.image_encoder_path\n",
        ").to(dtype=weight_dtype, device=\"cuda\")\n",
        "\n",
        "sched_kwargs = OmegaConf.to_container(infer_config.noise_scheduler_kwargs)\n",
        "scheduler = DDIMScheduler(**sched_kwargs)\n",
        "\n",
        "# load pretrained weights\n",
        "denoising_unet.load_state_dict(\n",
        "    torch.load(config.denoising_unet_path, map_location=\"cpu\"),\n",
        "    strict=False,\n",
        ")\n",
        "reference_unet.load_state_dict(\n",
        "    torch.load(config.reference_unet_path, map_location=\"cpu\"),\n",
        ")\n",
        "pose_guider.load_state_dict(\n",
        "    torch.load(config.pose_guider_path, map_location=\"cpu\"),\n",
        ")\n",
        "\n",
        "pipe = Pose2VideoPipeline(\n",
        "    vae=vae,\n",
        "    image_encoder=image_enc,\n",
        "    reference_unet=reference_unet,\n",
        "    denoising_unet=denoising_unet,\n",
        "    pose_guider=pose_guider,\n",
        "    scheduler=scheduler,\n",
        ")\n",
        "pipe = pipe.to(\"cuda\", dtype=weight_dtype)\n",
        "\n",
        "frame_inter_model = init_frame_interpolation_model()\n",
        "\n",
        "#def upscale_video_with_face_enhancer(video_path, method):\n",
        "    # Load the face enhancement model\n",
        "    #face_enhancer = load_face_enhancer_model(method)\n",
        "\n",
        "    # Upscale the video using the face enhancement model\n",
        "    #input_video = ffmpeg.input(video_path)\n",
        "    #output_video = ffmpeg.output(input_video.video, f\"{os.path.splitext(video_path)[0]}_upscaled.mp4\", vcodec='libx264', r=30)\n",
        "    #ffmpeg.run(output_video)\n",
        "\n",
        "   # return f\"{os.path.splitext(video_path)[0]}_upscaled.mp4\"\n",
        "\n",
        "\n",
        "\n",
        "####UPSCALE VIDEOS#####\n",
        "\n",
        "def upscale_video_with_face_enhancer(video_path, method, device, batch_size=12):\n",
        "    # Load the face enhancement model\n",
        "    face_enhancer, face_enhancer_runner = load_face_enhancer_model(method, device= 'cuda:0')\n",
        "\n",
        "    # Create a temporary directory to store the frames\n",
        "    with tempfile.TemporaryDirectory() as temp_dir:\n",
        "        # Extract frames from the input video\n",
        "        frame_paths = []\n",
        "        cap = cv2.VideoCapture(video_path)\n",
        "        frame_idx = 0\n",
        "        total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))\n",
        "        fps = cap.get(cv2.CAP_PROP_FPS)\n",
        "\n",
        "        pbar = tqdm(total=total_frames, unit=\"frames\")  # Create a progress bar\n",
        "\n",
        "        while True:\n",
        "            ret, frame = cap.read()\n",
        "            if not ret:\n",
        "                break\n",
        "            frame_path = os.path.join(temp_dir, f\"frame_{frame_idx}.jpg\")\n",
        "            cv2.imwrite(frame_path, frame)\n",
        "            frame_paths.append(frame_path)\n",
        "            frame_idx += 1\n",
        "            pbar.update(1)  # Update the progress bar\n",
        "\n",
        "        cap.release()\n",
        "        pbar.close()  # Close the progress bar\n",
        "\n",
        "        # Upscale the frames in batches\n",
        "        upscaled_frames = []\n",
        "        for i in tqdm(range(0, len(frame_paths), batch_size), unit=\"batches\", desc=\"Upscaling frames\"):\n",
        "            batch_frame_paths = frame_paths[i:i+batch_size]\n",
        "            batch_frames = [cv2.imread(path) for path in batch_frame_paths]\n",
        "            if face_enhancer is None:\n",
        "                batch_upscaled_frames = [face_enhancer_runner(frame, None) for frame in batch_frames]\n",
        "            else:\n",
        "                batch_upscaled_frames = [face_enhancer_runner(frame, face_enhancer) for frame in batch_frames]\n",
        "            upscaled_frames.extend(batch_upscaled_frames)\n",
        "\n",
        "        # Create the output video from the upscaled frames\n",
        "        with tempfile.NamedTemporaryFile(suffix=\".mp4\", delete=False) as temp_file:\n",
        "            fourcc = cv2.VideoWriter_fourcc(*'mp4v')\n",
        "            height, width, _ = upscaled_frames[0].shape\n",
        "            out = cv2.VideoWriter(temp_file.name, fourcc, fps, (width, height))\n",
        "            for frame in upscaled_frames:\n",
        "                out.write(frame)\n",
        "            out.release()\n",
        "\n",
        "            # Extract the audio from the original video\n",
        "            audio_output = f\"{temp_dir}/audio.aac\"\n",
        "            try:\n",
        "                ffmpeg.input(video_path).output(audio_output, acodec='copy').run()\n",
        "            except:\n",
        "                print(f\"Failed to extract audio from {video_path}\")\n",
        "                audio_output = None\n",
        "\n",
        "            # Generate the final video path with \"_upscaled\" at the end\n",
        "            unique_id = str(uuid.uuid4())\n",
        "            final_video_path = f\"{os.path.splitext(video_path)[0]}_{unique_id}_upscaled.mp4\"\n",
        "\n",
        "            # Merge the upscaled video and the audio\n",
        "            stream = ffmpeg.input(temp_file.name)\n",
        "            if audio_output:\n",
        "                audio = ffmpeg.input(audio_output)\n",
        "                ffmpeg.output(stream.video, audio.audio, final_video_path, vcodec='copy', acodec='aac', shortest=None).run()\n",
        "            else:\n",
        "                ffmpeg.output(stream.video, final_video_path, vcodec='copy').run()\n",
        "\n",
        "    # Return the path to the final video\n",
        "    return final_video_path\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "def get_headpose_temp(input_video):\n",
        "    lmk_extractor = LMKExtractor()\n",
        "    cap = cv2.VideoCapture(input_video)\n",
        "\n",
        "    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))\n",
        "    fps = cap.get(cv2.CAP_PROP_FPS)\n",
        "\n",
        "    trans_mat_list = []\n",
        "    while cap.isOpened():\n",
        "        ret, frame = cap.read()\n",
        "        if not ret:\n",
        "            break\n",
        "\n",
        "        result = lmk_extractor(frame)\n",
        "        trans_mat_list.append(result['trans_mat'].astype(np.float32))\n",
        "    cap.release()\n",
        "\n",
        "    trans_mat_arr = np.array(trans_mat_list)\n",
        "\n",
        "    # compute delta pose\n",
        "    trans_mat_inv_frame_0 = np.linalg.inv(trans_mat_arr[0])\n",
        "    pose_arr = np.zeros([trans_mat_arr.shape[0], 6])\n",
        "\n",
        "    for i in range(pose_arr.shape[0]):\n",
        "        pose_mat = trans_mat_inv_frame_0 @ trans_mat_arr[i]\n",
        "        euler_angles, translation_vector = matrix_to_euler_and_translation(pose_mat)\n",
        "        pose_arr[i, :3] =  euler_angles\n",
        "        pose_arr[i, 3:6] =  translation_vector\n",
        "\n",
        "    # interpolate to 30 fps\n",
        "    new_fps = 30\n",
        "    old_time = np.linspace(0, total_frames / fps, total_frames)\n",
        "    new_time = np.linspace(0, total_frames / fps, int(total_frames * new_fps / fps))\n",
        "\n",
        "    pose_arr_interp = np.zeros((len(new_time), 6))\n",
        "    for i in range(6):\n",
        "        interp_func = interp1d(old_time, pose_arr[:, i])\n",
        "        pose_arr_interp[:, i] = interp_func(new_time)\n",
        "\n",
        "    pose_arr_smooth = smooth_pose_seq(pose_arr_interp)\n",
        "\n",
        "    return pose_arr_smooth\n",
        "\n",
        "def audio2video(input_audio, ref_img, headpose_video=None, size=512, steps=25, length=60, seed=42, acc_flag=True):\n",
        "    fps = 30\n",
        "    cfg = 3.5\n",
        "    fi_step = 3 if acc_flag else 1\n",
        "\n",
        "    generator = torch.manual_seed(seed)\n",
        "\n",
        "    lmk_extractor = LMKExtractor()\n",
        "    vis = FaceMeshVisualizer()\n",
        "\n",
        "    width, height = size, size\n",
        "\n",
        "    date_str = datetime.now().strftime(\"%Y%m%d\")\n",
        "    time_str = datetime.now().strftime(\"%H%M\")\n",
        "    save_dir_name = f\"{time_str}--seed_{seed}-{size}x{size}\"\n",
        "\n",
        "    save_dir = Path(f\"output/{date_str}/{save_dir_name}\")\n",
        "    while os.path.exists(save_dir):\n",
        "        save_dir = Path(f\"output/{date_str}/{save_dir_name}_{np.random.randint(10000):04d}\")\n",
        "    save_dir.mkdir(exist_ok=True, parents=True)\n",
        "\n",
        "    ref_image_np = cv2.cvtColor(ref_img, cv2.COLOR_RGB2BGR)\n",
        "    ref_image_np = crop_face(ref_image_np, lmk_extractor)\n",
        "    if ref_image_np is None:\n",
        "        return None, Image.fromarray(ref_img)\n",
        "\n",
        "    ref_image_np = cv2.resize(ref_image_np, (size, size))\n",
        "    ref_image_pil = Image.fromarray(cv2.cvtColor(ref_image_np, cv2.COLOR_BGR2RGB))\n",
        "\n",
        "    face_result = lmk_extractor(ref_image_np)\n",
        "    if face_result is None:\n",
        "        return None, ref_image_pil\n",
        "\n",
        "    lmks = face_result['lmks'].astype(np.float32)\n",
        "    ref_pose = vis.draw_landmarks((ref_image_np.shape[1], ref_image_np.shape[0]), lmks, normed=True)\n",
        "\n",
        "    sample = prepare_audio_feature(input_audio, wav2vec_model_path=audio_infer_config['a2m_model']['model_path'])\n",
        "    sample['audio_feature'] = torch.from_numpy(sample['audio_feature']).float().cuda()\n",
        "    sample['audio_feature'] = sample['audio_feature'].unsqueeze(0)\n",
        "\n",
        "    # inference\n",
        "    pred = a2m_model.infer(sample['audio_feature'], sample['seq_len'])\n",
        "    pred = pred.squeeze().detach().cpu().numpy()\n",
        "    pred = pred.reshape(pred.shape[0], -1, 3)\n",
        "    pred = pred + face_result['lmks3d']\n",
        "\n",
        "    if headpose_video is not None:\n",
        "        pose_seq = get_headpose_temp(headpose_video)\n",
        "    else:\n",
        "        pose_seq = np.load(config['pose_temp'])\n",
        "    mirrored_pose_seq = np.concatenate((pose_seq, pose_seq[-2:0:-1]), axis=0)\n",
        "    cycled_pose_seq = np.tile(mirrored_pose_seq, (sample['seq_len'] // len(mirrored_pose_seq) + 1, 1))[:sample['seq_len']]\n",
        "\n",
        "    # project 3D mesh to 2D landmark\n",
        "    projected_vertices = project_points(pred, face_result['trans_mat'], cycled_pose_seq, [height, width])\n",
        "\n",
        "    pose_images = []\n",
        "    for i, verts in enumerate(projected_vertices):\n",
        "        lmk_img = vis.draw_landmarks((width, height), verts, normed=False)\n",
        "        pose_images.append(lmk_img)\n",
        "\n",
        "    pose_list = []\n",
        "    args_L = len(pose_images) if length==0 or length > len(pose_images) else length\n",
        "    for pose_image_np in pose_images[: args_L : fi_step]:\n",
        "        pose_image_np = cv2.resize(pose_image_np,  (width, height))\n",
        "        pose_list.append(pose_image_np)\n",
        "\n",
        "    pose_list = np.array(pose_list)\n",
        "\n",
        "    video_length = len(pose_list)\n",
        "\n",
        "    video = pipe(\n",
        "        ref_image_pil,\n",
        "        pose_list,\n",
        "        ref_pose,\n",
        "        width,\n",
        "        height,\n",
        "        video_length,\n",
        "        steps,\n",
        "        cfg,\n",
        "        generator=generator,\n",
        "    ).videos\n",
        "\n",
        "    if acc_flag:\n",
        "        video = batch_images_interpolation_tool(video, frame_inter_model, inter_frames=fi_step-1)\n",
        "\n",
        "    save_path = f\"{save_dir}/{size}x{size}_{time_str}_noaudio.mp4\"\n",
        "    save_videos_grid(\n",
        "        video,\n",
        "        save_path,\n",
        "        n_rows=1,\n",
        "        fps=fps,\n",
        "    )\n",
        "\n",
        "    stream = ffmpeg.input(save_path)\n",
        "    audio = ffmpeg.input(input_audio)\n",
        "    ffmpeg.output(stream.video, audio.audio, save_path.replace('_noaudio.mp4', '.mp4'), vcodec='copy', acodec='aac', shortest=None).run()\n",
        "    os.remove(save_path)\n",
        "\n",
        "    return save_path.replace('_noaudio.mp4', '.mp4'), ref_image_pil\n",
        "\n",
        "def video2video(ref_img, source_video, size=512, steps=25, length=60, seed=42, acc_flag=True):\n",
        "    cfg = 3.5\n",
        "    fi_step = 3 if acc_flag else 1\n",
        "\n",
        "    generator = torch.manual_seed(seed)\n",
        "\n",
        "    lmk_extractor = LMKExtractor()\n",
        "    vis = FaceMeshVisualizer()\n",
        "\n",
        "    width, height = size, size\n",
        "\n",
        "    date_str = datetime.now().strftime(\"%Y%m%d\")\n",
        "    time_str = datetime.now().strftime(\"%H%M\")\n",
        "    save_dir_name = f\"{time_str}--seed_{seed}-{size}x{size}\"\n",
        "\n",
        "    save_dir = Path(f\"output/{date_str}/{save_dir_name}\")\n",
        "    while os.path.exists(save_dir):\n",
        "        save_dir = Path(f\"output/{date_str}/{save_dir_name}_{np.random.randint(10000):04d}\")\n",
        "    save_dir.mkdir(exist_ok=True, parents=True)\n",
        "\n",
        "    ref_image_np = cv2.cvtColor(ref_img, cv2.COLOR_RGB2BGR)\n",
        "    ref_image_np = crop_face(ref_image_np, lmk_extractor)\n",
        "    if ref_image_np is None:\n",
        "        return None, Image.fromarray(ref_img)\n",
        "\n",
        "    ref_image_np = cv2.resize(ref_image_np, (size, size))\n",
        "    ref_image_pil = Image.fromarray(cv2.cvtColor(ref_image_np, cv2.COLOR_BGR2RGB))\n",
        "\n",
        "    face_result = lmk_extractor(ref_image_np)\n",
        "    if face_result is None:\n",
        "        return None, ref_image_pil\n",
        "\n",
        "    lmks = face_result['lmks'].astype(np.float32)\n",
        "    ref_pose = vis.draw_landmarks((ref_image_np.shape[1], ref_image_np.shape[0]), lmks, normed=True)\n",
        "\n",
        "    source_images = read_frames(source_video)\n",
        "    src_fps = get_fps(source_video)\n",
        "\n",
        "    step = 1\n",
        "    if src_fps == 60:\n",
        "        src_fps = 30\n",
        "        step = 2\n",
        "\n",
        "    pose_trans_list = []\n",
        "    verts_list = []\n",
        "    bs_list = []\n",
        "    args_L = len(source_images) if length==0 or length*step > len(source_images) else length*step\n",
        "    for src_image_pil in source_images[: args_L : step*fi_step]:\n",
        "        src_img_np = cv2.cvtColor(np.array(src_image_pil), cv2.COLOR_RGB2BGR)\n",
        "        frame_height, frame_width, _ = src_img_np.shape\n",
        "        src_img_result = lmk_extractor(src_img_np)\n",
        "        if src_img_result is None:\n",
        "            break\n",
        "        pose_trans_list.append(src_img_result['trans_mat'])\n",
        "        verts_list.append(src_img_result['lmks3d'])\n",
        "        bs_list.append(src_img_result['bs'])\n",
        "\n",
        "    trans_mat_arr = np.array(pose_trans_list)\n",
        "    verts_arr = np.array(verts_list)\n",
        "    bs_arr = np.array(bs_list)\n",
        "    min_bs_idx = np.argmin(bs_arr.sum(1))\n",
        "\n",
        "    # compute delta pose\n",
        "    pose_arr = np.zeros([trans_mat_arr.shape[0], 6])\n",
        "\n",
        "    for i in range(pose_arr.shape[0]):\n",
        "        euler_angles, translation_vector = matrix_to_euler_and_translation(trans_mat_arr[i]) # real pose of source\n",
        "        pose_arr[i, :3] =  euler_angles\n",
        "        pose_arr[i, 3:6] =  translation_vector\n",
        "\n",
        "    init_tran_vec = face_result['trans_mat'][:3, 3] # init translation of tgt\n",
        "    pose_arr[:, 3:6] = pose_arr[:, 3:6] - pose_arr[0, 3:6] + init_tran_vec # (relative translation of source) + (init translation of tgt)\n",
        "\n",
        "    pose_arr_smooth = smooth_pose_seq(pose_arr, window_size=3)\n",
        "    pose_mat_smooth = [euler_and_translation_to_matrix(pose_arr_smooth[i][:3], pose_arr_smooth[i][3:6]) for i in range(pose_arr_smooth.shape[0])]\n",
        "    pose_mat_smooth = np.array(pose_mat_smooth)\n",
        "\n",
        "    # face retarget\n",
        "    verts_arr = verts_arr - verts_arr[min_bs_idx] + face_result['lmks3d']\n",
        "    # project 3D mesh to 2D landmark\n",
        "    projected_vertices = project_points_with_trans(verts_arr, pose_mat_smooth, [frame_height, frame_width])\n",
        "\n",
        "    pose_list = []\n",
        "    for i, verts in enumerate(projected_vertices):\n",
        "        lmk_img = vis.draw_landmarks((frame_width, frame_height), verts, normed=False)\n",
        "        pose_image_np = cv2.resize(lmk_img,  (width, height))\n",
        "        pose_list.append(pose_image_np)\n",
        "\n",
        "    pose_list = np.array(pose_list)\n",
        "\n",
        "    video_length = len(pose_list)\n",
        "\n",
        "    video = pipe(\n",
        "        ref_image_pil,\n",
        "        pose_list,\n",
        "        ref_pose,\n",
        "        width,\n",
        "        height,\n",
        "        video_length,\n",
        "        steps,\n",
        "        cfg,\n",
        "        generator=generator,\n",
        "    ).videos\n",
        "\n",
        "    if acc_flag:\n",
        "        video = batch_images_interpolation_tool(video, frame_inter_model, inter_frames=fi_step-1)\n",
        "\n",
        "    save_path = f\"{save_dir}/{size}x{size}_{time_str}_noaudio.mp4\"\n",
        "    save_videos_grid(\n",
        "        video,\n",
        "        save_path,\n",
        "        n_rows=1,\n",
        "        fps=src_fps,\n",
        "    )\n",
        "\n",
        "    audio_output = f'{save_dir}/audio_from_video.aac'\n",
        "    # extract audio\n",
        "    try:\n",
        "        ffmpeg.input(source_video).output(audio_output, acodec='copy').run()\n",
        "        # merge audio and video\n",
        "        stream = ffmpeg.input(save_path)\n",
        "        audio = ffmpeg.input(audio_output)\n",
        "        ffmpeg.output(stream.video, audio.audio, save_path.replace('_noaudio.mp4', '.mp4'), vcodec='copy', acodec='aac', shortest=None).run()\n",
        "\n",
        "        os.remove(save_path)\n",
        "        os.remove(audio_output)\n",
        "    except:\n",
        "        shutil.move(\n",
        "            save_path,\n",
        "            save_path.replace('_noaudio.mp4', '.mp4')\n",
        "        )\n",
        "\n",
        "    return save_path.replace('_noaudio.mp4', '.mp4'), ref_image_pil\n",
        "\n",
        "\n",
        "################# GUI ################\n",
        "\n",
        "title = r\"\"\"\n",
        "<h1>AniPortrait</h1>\n",
        "\"\"\"\n",
        "\n",
        "description = r\"\"\"\n",
        "<b>Official 🤗 Gradio demo</b> for <a href='https://github.com/Zejun-Yang/AniPortrait' target='_blank'><b>AniPortrait: Audio-Driven Synthesis of Photorealistic Portrait Animations</b></a>.<br>\n",
        "\"\"\"\n",
        "\n",
        "with gr.Blocks() as demo:\n",
        "\n",
        "    gr.Markdown(title)\n",
        "    gr.Markdown(description)\n",
        "\n",
        "    with gr.Tab(\"Audio2video\"):\n",
        "        with gr.Row():\n",
        "            with gr.Column():\n",
        "                with gr.Row():\n",
        "                    a2v_input_audio = gr.Audio(sources=[\"upload\", \"microphone\"], type=\"filepath\", editable=True, label=\"Input audio\", interactive=True)\n",
        "                    a2v_ref_img = gr.Image(label=\"Upload reference image\", sources=\"upload\")\n",
        "                    a2v_headpose_video = gr.Video(label=\"Option: upload head pose reference video\", sources=\"upload\")\n",
        "\n",
        "                with gr.Row():\n",
        "                    a2v_size_slider = gr.Slider(minimum=256, maximum=768, step=8, value=512, label=\"Video size (-W & -H)\")\n",
        "                    a2v_step_slider = gr.Slider(minimum=5, maximum=30, step=1, value=25, label=\"Steps (--steps)\")\n",
        "\n",
        "                with gr.Row():\n",
        "                    a2v_length = gr.Slider(minimum=0, maximum=9999, step=1, value=60, label=\"Length (-L) (Set to 0 to automatically calculate length)\")\n",
        "                    a2v_seed = gr.Number(value=42, label=\"Seed (--seed)\")\n",
        "\n",
        "                with gr.Row():\n",
        "                    a2v_acc_flag = gr.Checkbox(value=True, label=\"Accelerate\")\n",
        "                    a2v_botton = gr.Button(\"Generate\", variant=\"primary\")\n",
        "            a2v_output_video = gr.PlayableVideo(label=\"Result\", interactive=False)\n",
        "\n",
        "        #gr.Examples(\n",
        "          #@  examples=[\n",
        "          #      [\"configs/inference/audio/lyl.wav\", \"configs/inference/ref_images/Aragaki.png\", None],\n",
        "          #      [\"configs/inference/audio/lyl.wav\", \"configs/inference/ref_images/solo.png\", None],\n",
        "           #     [\"configs/inference/audio/lyl.wav\", \"configs/inference/ref_images/lyl.png\", \"configs/inference/head_pose_temp/pose_ref_video.mp4\"],\n",
        "           #     ],\n",
        "          #  inputs=[a2v_input_audio, a2v_ref_img, a2v_headpose_video],\n",
        "        #)\n",
        "\n",
        "    with gr.Tab(\"Video2video\"):\n",
        "        with gr.Row():\n",
        "            with gr.Column():\n",
        "                with gr.Row():\n",
        "                    v2v_ref_img = gr.Image(label=\"Upload reference image\", sources=\"upload\")\n",
        "                    v2v_source_video = gr.Video(label=\"Upload source video\", sources=\"upload\")\n",
        "\n",
        "                with gr.Row():\n",
        "                    v2v_size_slider = gr.Slider(minimum=256, maximum=768, step=8, value=512, label=\"Video size (-W & -H)\")\n",
        "                    v2v_step_slider = gr.Slider(minimum=5, maximum=30, step=1, value=25, label=\"Steps (--steps)\")\n",
        "\n",
        "                with gr.Row():\n",
        "                    v2v_length = gr.Slider(minimum=0, maximum=9999, step=1, value=60, label=\"Length (-L) (Set to 0 to automatically calculate length)\")\n",
        "                    v2v_seed = gr.Number(value=42, label=\"Seed (--seed)\")\n",
        "\n",
        "                with gr.Row():\n",
        "                    v2v_acc_flag = gr.Checkbox(value=True, label=\"Accelerate\")\n",
        "                    v2v_botton = gr.Button(\"Generate\", variant=\"primary\")\n",
        "            v2v_output_video = gr.PlayableVideo(label=\"Result\", interactive=False)\n",
        "\n",
        "        #gr.Examples(\n",
        "         #   examples=[\n",
        "         #       [\"configs/inference/ref_images/Aragaki.png\", \"configs/inference/video/Aragaki_song.mp4\"],\n",
        "          #      [\"configs/inference/ref_images/solo.png\", \"configs/inference/video/Aragaki_song.mp4\"],\n",
        "          #      [\"configs/inference/ref_images/lyl.png\", \"configs/inference/head_pose_temp/pose_ref_video.mp4\"],\n",
        "          #      ],\n",
        "          #  inputs=[v2v_ref_img, v2v_source_video, a2v_headpose_video],\n",
        "        #)\n",
        "    with gr.Tab(\"Video Upscale\"):\n",
        "        with gr.Row():\n",
        "            with gr.Column():\n",
        "                with gr.Row():\n",
        "                  video_path_textbox = gr.Textbox(label=\"Enter video path:\")\n",
        "                  upscale_method = gr.Dropdown(\n",
        "                    get_available_enhancer_names(),\n",
        "                    label=\"Upscale method\",\n",
        "                    value=\"REAL-ESRGAN 4x\",\n",
        "                )\n",
        "        upscale_botton = gr.Button(\"Upscale\", variant=\"primary\")\n",
        "        upscale_output_video = gr.PlayableVideo(\n",
        "        label=\"Upscaled video\", interactive=False\n",
        "        )\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "    a2v_botton.click(\n",
        "        fn=audio2video,\n",
        "        inputs=[a2v_input_audio, a2v_ref_img, a2v_headpose_video,\n",
        "                a2v_size_slider, a2v_step_slider, a2v_length, a2v_seed, a2v_acc_flag],\n",
        "        outputs=[a2v_output_video, a2v_ref_img]\n",
        "    )\n",
        "    v2v_botton.click(\n",
        "        fn=video2video,\n",
        "        inputs=[v2v_ref_img, v2v_source_video,\n",
        "                v2v_size_slider, v2v_step_slider, v2v_length, v2v_seed, v2v_acc_flag],\n",
        "        outputs=[v2v_output_video, v2v_ref_img]\n",
        "    )\n",
        "\n",
        "    upscale_botton.click(\n",
        "        fn=upscale_video_with_face_enhancer,\n",
        "        inputs=[video_path_textbox, upscale_method],\n",
        "        outputs=[upscale_output_video],\n",
        ")\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "demo.launch(share=True)\n"
      ],
      "metadata": {
        "cellView": "form",
        "id": "8OYQmbCh3Zys"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title UPSCALE VIDEO\n",
        "import tempfile\n",
        "import uuid\n",
        "import os\n",
        "import cv2\n",
        "import ffmpeg\n",
        "from face_enhancer import load_face_enhancer_model\n",
        "from tqdm import tqdm\n",
        "\n",
        "def upscale_video_with_face_enhancer(video_path, method, device, batch_size=8):\n",
        "    # Load the face enhancement model\n",
        "    face_enhancer, face_enhancer_runner = load_face_enhancer_model(method, device=device)\n",
        "\n",
        "    # Create a temporary directory to store the frames\n",
        "    with tempfile.TemporaryDirectory() as temp_dir:\n",
        "        # Extract frames from the input video\n",
        "        frame_paths = []\n",
        "        cap = cv2.VideoCapture(video_path)\n",
        "        frame_idx = 0\n",
        "        total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))\n",
        "        fps = cap.get(cv2.CAP_PROP_FPS)\n",
        "\n",
        "        pbar = tqdm(total=total_frames, unit=\"frames\")  # Create a progress bar\n",
        "\n",
        "        while True:\n",
        "            ret, frame = cap.read()\n",
        "            if not ret:\n",
        "                break\n",
        "            frame_path = os.path.join(temp_dir, f\"frame_{frame_idx}.jpg\")\n",
        "            cv2.imwrite(frame_path, frame)\n",
        "            frame_paths.append(frame_path)\n",
        "            frame_idx += 1\n",
        "            pbar.update(1)  # Update the progress bar\n",
        "\n",
        "        cap.release()\n",
        "        pbar.close()  # Close the progress bar\n",
        "\n",
        "        # Upscale the frames in batches\n",
        "        upscaled_frames = []\n",
        "        for i in tqdm(range(0, len(frame_paths), batch_size), unit=\"batches\", desc=\"Upscaling frames\"):\n",
        "            batch_frame_paths = frame_paths[i:i+batch_size]\n",
        "            batch_frames = [cv2.imread(path) for path in batch_frame_paths]\n",
        "            if face_enhancer is None:\n",
        "                batch_upscaled_frames = [face_enhancer_runner(frame, None) for frame in batch_frames]\n",
        "            else:\n",
        "                batch_upscaled_frames = [face_enhancer_runner(frame, face_enhancer) for frame in batch_frames]\n",
        "            upscaled_frames.extend(batch_upscaled_frames)\n",
        "\n",
        "        # Create the output video from the upscaled frames\n",
        "        with tempfile.NamedTemporaryFile(suffix=\".mp4\", delete=True) as temp_file:\n",
        "            fourcc = cv2.VideoWriter_fourcc(*'mp4v')\n",
        "            height, width, _ = upscaled_frames[0].shape\n",
        "            out = cv2.VideoWriter(temp_file.name, fourcc, fps, (width, height))\n",
        "            for frame in upscaled_frames:\n",
        "                out.write(frame)\n",
        "            out.release()\n",
        "\n",
        "            # Extract the audio from the original video\n",
        "            audio_output = f\"{temp_dir}/audio.aac\"\n",
        "            try:\n",
        "                ffmpeg.input(video_path).output(audio_output, acodec='copy').run()\n",
        "            except:\n",
        "                print(f\"Failed to extract audio from {video_path}\")\n",
        "                audio_output = None\n",
        "\n",
        "            # Generate the final video path with \"_upscaled\" at the end\n",
        "            unique_id = str(uuid.uuid4())\n",
        "            final_video_path = f\"{os.path.splitext(video_path)[0]}_{unique_id}_upscaled.mp4\"\n",
        "\n",
        "            # Merge the upscaled video and the audio\n",
        "            stream = ffmpeg.input(temp_file.name)\n",
        "            if audio_output:\n",
        "                audio = ffmpeg.input(audio_output)\n",
        "                ffmpeg.output(stream.video, audio.audio, final_video_path, vcodec='copy', acodec='aac', shortest=None).run()\n",
        "            else:\n",
        "                ffmpeg.output(stream.video, final_video_path, vcodec='copy').run()\n",
        "\n",
        "    # Return the path to the final video\n",
        "    return final_video_path"
      ],
      "metadata": {
        "cellView": "form",
        "id": "eyKiAOby2NOc"
      },
      "execution_count": null,
      "outputs": []
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "gpuType": "T4",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}