api.py

# -*- coding: utf-8 -*-
# @Time    : 2024/9/13 0:23
# @Project : FasterLivePortrait
# @FileName: api.py
import pdb
import shutil
from typing import Optional, Dict, Any
import io
import os
import subprocess
import uvicorn
import cv2
import time
import numpy as np
import os
import datetime
import platform
import pickle
from tqdm import tqdm
from pydantic import BaseModel
from fastapi import APIRouter, Depends, FastAPI, Request, Response, UploadFile
from fastapi import File, Body, Form
from omegaconf import OmegaConf
from fastapi.responses import StreamingResponse
from zipfile import ZipFile
from src.pipelines.faster_live_portrait_pipeline import FasterLivePortraitPipeline
from src.utils.utils import video_has_audio
from src.utils import logger

# model dir
project_dir = os.path.dirname(__file__)
checkpoints_dir = os.environ.get("FLIP_CHECKPOINT_DIR", os.path.join(project_dir, "checkpoints"))
log_dir = os.path.join(project_dir, "logs")
os.makedirs(log_dir, exist_ok=True)
result_dir = os.path.join(project_dir, "results")
os.makedirs(result_dir, exist_ok=True)

logger_f = logger.get_logger("faster_liveportrait_api", log_file=os.path.join(log_dir, "log_run.log"))

app = FastAPI()

global pipe

if platform.system().lower() == 'windows':
    FFMPEG = "third_party/ffmpeg-7.0.1-full_build/bin/ffmpeg.exe"
else:
    FFMPEG = "ffmpeg"


def check_all_checkpoints_exist(infer_cfg):
    """
    check whether all checkpoints exist
    :return:
    """
    ret = True
    for name in infer_cfg.models:
        if not isinstance(infer_cfg.models[name].model_path, str):
            for i in range(len(infer_cfg.models[name].model_path)):
                infer_cfg.models[name].model_path[i] = infer_cfg.models[name].model_path[i].replace("./checkpoints",
                                                                                                    checkpoints_dir)
                if not os.path.exists(infer_cfg.models[name].model_path[i]) and not os.path.exists(
                        infer_cfg.models[name].model_path[i][:-4] + ".onnx"):
                    return False
        else:
            infer_cfg.models[name].model_path = infer_cfg.models[name].model_path.replace("./checkpoints",
                                                                                          checkpoints_dir)
            if not os.path.exists(infer_cfg.models[name].model_path) and not os.path.exists(
                    infer_cfg.models[name].model_path[:-4] + ".onnx"):
                return False
    for name in infer_cfg.animal_models:
        if not isinstance(infer_cfg.animal_models[name].model_path, str):
            for i in range(len(infer_cfg.animal_models[name].model_path)):
                infer_cfg.animal_models[name].model_path[i] = infer_cfg.animal_models[name].model_path[i].replace(
                    "./checkpoints",
                    checkpoints_dir)
                if not os.path.exists(infer_cfg.animal_models[name].model_path[i]) and not os.path.exists(
                        infer_cfg.animal_models[name].model_path[i][:-4] + ".onnx"):
                    return False
        else:
            infer_cfg.animal_models[name].model_path = infer_cfg.animal_models[name].model_path.replace("./checkpoints",
                                                                                                        checkpoints_dir)
            if not os.path.exists(infer_cfg.animal_models[name].model_path) and not os.path.exists(
                    infer_cfg.animal_models[name].model_path[:-4] + ".onnx"):
                return False

    # XPOSE
    xpose_model_path = os.path.join(checkpoints_dir, "liveportrait_animal_onnx/xpose.pth")
    if not os.path.exists(xpose_model_path):
        return False
    embeddings_cache_9_path = os.path.join(checkpoints_dir, "liveportrait_animal_onnx/clip_embedding_9.pkl")
    if not os.path.exists(embeddings_cache_9_path):
        return False
    embeddings_cache_68_path = os.path.join(checkpoints_dir, "liveportrait_animal_onnx/clip_embedding_68.pkl")
    if not os.path.exists(embeddings_cache_68_path):
        return False
    return ret


def convert_onnx_to_trt_models(infer_cfg):
    ret = True
    for name in infer_cfg.models:
        if not isinstance(infer_cfg.models[name].model_path, str):
            for i in range(len(infer_cfg.models[name].model_path)):
                trt_path = infer_cfg.models[name].model_path[i]
                onnx_path = trt_path[:-4] + ".onnx"
                if not os.path.exists(trt_path):
                    convert_cmd = f"python scripts/onnx2trt.py -o {onnx_path}"
                    logger_f.info(f"convert onnx model: {onnx_path}")
                    result = subprocess.run(convert_cmd, shell=True, check=True)
                    # 检查结果
                    if result.returncode == 0:
                        logger_f.info(f"convert onnx model: {onnx_path} successful")
                    else:
                        logger_f.error(f"convert onnx model: {onnx_path} failed")
                        return False
        else:
            trt_path = infer_cfg.models[name].model_path
            onnx_path = trt_path[:-4] + ".onnx"
            if not os.path.exists(trt_path):
                convert_cmd = f"python scripts/onnx2trt.py -o {onnx_path}"
                logger_f.info(f"convert onnx model: {onnx_path}")
                result = subprocess.run(convert_cmd, shell=True, check=True)
                # 检查结果
                if result.returncode == 0:
                    logger_f.info(f"convert onnx model: {onnx_path} successful")
                else:
                    logger_f.error(f"convert onnx model: {onnx_path} failed")
                    return False

    for name in infer_cfg.animal_models:
        if not isinstance(infer_cfg.animal_models[name].model_path, str):
            for i in range(len(infer_cfg.animal_models[name].model_path)):
                trt_path = infer_cfg.animal_models[name].model_path[i]
                onnx_path = trt_path[:-4] + ".onnx"
                if not os.path.exists(trt_path):
                    convert_cmd = f"python scripts/onnx2trt.py -o {onnx_path}"
                    logger_f.info(f"convert onnx model: {onnx_path}")
                    result = subprocess.run(convert_cmd, shell=True, check=True)
                    # 检查结果
                    if result.returncode == 0:
                        logger_f.info(f"convert onnx model: {onnx_path} successful")
                    else:
                        logger_f.error(f"convert onnx model: {onnx_path} failed")
                        return False
        else:
            trt_path = infer_cfg.animal_models[name].model_path
            onnx_path = trt_path[:-4] + ".onnx"
            if not os.path.exists(trt_path):
                convert_cmd = f"python scripts/onnx2trt.py -o {onnx_path}"
                logger_f.info(f"convert onnx model: {onnx_path}")
                result = subprocess.run(convert_cmd, shell=True, check=True)
                # 检查结果
                if result.returncode == 0:
                    logger_f.info(f"convert onnx model: {onnx_path} successful")
                else:
                    logger_f.error(f"convert onnx model: {onnx_path} failed")
                    return False
    return ret


@app.on_event("startup")
async def startup_event():
    global pipe
    # default use trt model
    cfg_file = os.path.join(project_dir, "configs/trt_infer.yaml")
    infer_cfg = OmegaConf.load(cfg_file)
    checkpoints_exist = check_all_checkpoints_exist(infer_cfg)

    # first: download model if not exist
    if not checkpoints_exist:
        download_cmd = f"huggingface-cli download warmshao/FasterLivePortrait --local-dir {checkpoints_dir}"
        logger_f.info(f"download model: {download_cmd}")
        result = subprocess.run(download_cmd, shell=True, check=True)
        # 检查结果
        if result.returncode == 0:
            logger_f.info(f"Download checkpoints to {checkpoints_dir} successful")
        else:
            logger_f.error(f"Download checkpoints to {checkpoints_dir} failed")
            exit(1)
    # second: convert onnx model to trt
    convert_ret = convert_onnx_to_trt_models(infer_cfg)
    if not convert_ret:
        logger_f.error(f"convert onnx model to trt failed")
        exit(1)

    infer_cfg.infer_params.flag_pasteback = True
    pipe = FasterLivePortraitPipeline(cfg=infer_cfg, is_animal=True)


def run_with_video(source_image_path, driving_video_path, save_dir):
    global pipe
    ret = pipe.prepare_source(source_image_path, realtime=False)
    if not ret:
        logger_f.warning(f"no face in {source_image_path}! exit!")
        return
    vcap = cv2.VideoCapture(driving_video_path)
    fps = int(vcap.get(cv2.CAP_PROP_FPS))
    h, w = pipe.src_imgs[0].shape[:2]

    # render output video
    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    vsave_crop_path = os.path.join(save_dir,
                                   f"{os.path.basename(source_image_path)}-{os.path.basename(driving_video_path)}-crop.mp4")
    vout_crop = cv2.VideoWriter(vsave_crop_path, fourcc, fps, (512 * 2, 512))
    vsave_org_path = os.path.join(save_dir,
                                  f"{os.path.basename(source_image_path)}-{os.path.basename(driving_video_path)}-org.mp4")
    vout_org = cv2.VideoWriter(vsave_org_path, fourcc, fps, (w, h))

    infer_times = []
    motion_lst = []
    c_eyes_lst = []
    c_lip_lst = []

    frame_ind = 0
    while vcap.isOpened():
        ret, frame = vcap.read()
        if not ret:
            break
        t0 = time.time()
        first_frame = frame_ind == 0
        dri_crop, out_crop, out_org, dri_motion_info = pipe.run(frame, pipe.src_imgs[0], pipe.src_infos[0],
                                                                first_frame=first_frame)
        frame_ind += 1
        if out_crop is None:
            logger_f.warning(f"no face in driving frame:{frame_ind}")
            continue

        motion_lst.append(dri_motion_info[0])
        c_eyes_lst.append(dri_motion_info[1])
        c_lip_lst.append(dri_motion_info[2])

        infer_times.append(time.time() - t0)
        # print(time.time() - t0)
        dri_crop = cv2.resize(dri_crop, (512, 512))
        out_crop = np.concatenate([dri_crop, out_crop], axis=1)
        out_crop = cv2.cvtColor(out_crop, cv2.COLOR_RGB2BGR)
        vout_crop.write(out_crop)
        out_org = cv2.cvtColor(out_org, cv2.COLOR_RGB2BGR)
        vout_org.write(out_org)
    vcap.release()
    vout_crop.release()
    vout_org.release()
    if video_has_audio(driving_video_path):
        vsave_crop_path_new = os.path.splitext(vsave_crop_path)[0] + "-audio.mp4"
        subprocess.call(
            [FFMPEG, "-i", vsave_crop_path, "-i", driving_video_path,
             "-b:v", "10M", "-c:v",
             "libx264", "-map", "0:v", "-map", "1:a",
             "-c:a", "aac",
             "-pix_fmt", "yuv420p", vsave_crop_path_new, "-y", "-shortest"])
        vsave_org_path_new = os.path.splitext(vsave_org_path)[0] + "-audio.mp4"
        subprocess.call(
            [FFMPEG, "-i", vsave_org_path, "-i", driving_video_path,
             "-b:v", "10M", "-c:v",
             "libx264", "-map", "0:v", "-map", "1:a",
             "-c:a", "aac",
             "-pix_fmt", "yuv420p", vsave_org_path_new, "-y", "-shortest"])

        logger_f.info(vsave_crop_path_new)
        logger_f.info(vsave_org_path_new)
    else:
        logger_f.info(vsave_crop_path)
        logger_f.info(vsave_org_path)

    logger_f.info(
        "inference median time: {} ms/frame, mean time: {} ms/frame".format(np.median(infer_times) * 1000,
                                                                            np.mean(infer_times) * 1000))
    # save driving motion to pkl
    template_dct = {
        'n_frames': len(motion_lst),
        'output_fps': fps,
        'motion': motion_lst,
        'c_eyes_lst': c_eyes_lst,
        'c_lip_lst': c_lip_lst,
    }
    template_pkl_path = os.path.join(save_dir,
                                     f"{os.path.basename(driving_video_path)}.pkl")
    with open(template_pkl_path, "wb") as fw:
        pickle.dump(template_dct, fw)
    logger_f.info(f"save driving motion pkl file at : {template_pkl_path}")


def run_with_pkl(source_image_path, driving_pickle_path, save_dir):
    global pipe
    ret = pipe.prepare_source(source_image_path, realtime=False)
    if not ret:
        logger_f.warning(f"no face in {source_image_path}! exit!")
        return

    with open(driving_pickle_path, "rb") as fin:
        dri_motion_infos = pickle.load(fin)

    fps = int(dri_motion_infos["output_fps"])
    h, w = pipe.src_imgs[0].shape[:2]

    # render output video
    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    vsave_crop_path = os.path.join(save_dir,
                                   f"{os.path.basename(source_image_path)}-{os.path.basename(driving_pickle_path)}-crop.mp4")
    vout_crop = cv2.VideoWriter(vsave_crop_path, fourcc, fps, (512, 512))
    vsave_org_path = os.path.join(save_dir,
                                  f"{os.path.basename(source_image_path)}-{os.path.basename(driving_pickle_path)}-org.mp4")
    vout_org = cv2.VideoWriter(vsave_org_path, fourcc, fps, (w, h))

    infer_times = []
    motion_lst = dri_motion_infos["motion"]
    c_eyes_lst = dri_motion_infos["c_eyes_lst"] if "c_eyes_lst" in dri_motion_infos else dri_motion_infos[
        "c_d_eyes_lst"]
    c_lip_lst = dri_motion_infos["c_lip_lst"] if "c_lip_lst" in dri_motion_infos else dri_motion_infos["c_d_lip_lst"]

    frame_num = len(motion_lst)
    for frame_ind in tqdm(range(frame_num)):
        t0 = time.time()
        first_frame = frame_ind == 0
        dri_motion_info_ = [motion_lst[frame_ind], c_eyes_lst[frame_ind], c_lip_lst[frame_ind]]
        out_crop, out_org = pipe.run_with_pkl(dri_motion_info_, pipe.src_imgs[0], pipe.src_infos[0],
                                              first_frame=first_frame)
        if out_crop is None:
            logger_f.warning(f"no face in driving frame:{frame_ind}")
            continue

        infer_times.append(time.time() - t0)
        # print(time.time() - t0)
        out_crop = cv2.cvtColor(out_crop, cv2.COLOR_RGB2BGR)
        vout_crop.write(out_crop)
        out_org = cv2.cvtColor(out_org, cv2.COLOR_RGB2BGR)
        vout_org.write(out_org)

    vout_crop.release()
    vout_org.release()
    logger_f.info(vsave_crop_path)
    logger_f.info(vsave_org_path)
    logger_f.info(
        "inference median time: {} ms/frame, mean time: {} ms/frame".format(np.median(infer_times) * 1000,
                                                                            np.mean(infer_times) * 1000))


class LivePortraitParams(BaseModel):
    flag_pickle: bool = False
    flag_relative_input: bool = True
    flag_do_crop_input: bool = True
    flag_remap_input: bool = True
    driving_multiplier: float = 1.0
    flag_stitching: bool = True
    flag_crop_driving_video_input: bool = True
    flag_video_editing_head_rotation: bool = False
    flag_is_animal: bool = True
    scale: float = 2.3
    vx_ratio: float = 0.0
    vy_ratio: float = -0.125
    scale_crop_driving_video: float = 2.2
    vx_ratio_crop_driving_video: float = 0.0
    vy_ratio_crop_driving_video: float = -0.1
    driving_smooth_observation_variance: float = 1e-7


@app.post("/predict/")
async def upload_files(
        source_image: Optional[UploadFile] = File(None),
        driving_video: Optional[UploadFile] = File(None),
        driving_pickle: Optional[UploadFile] = File(None),
        flag_is_animal: bool = Form(...),
        flag_pickle: bool = Form(...),
        flag_relative_input: bool = Form(...),
        flag_do_crop_input: bool = Form(...),
        flag_remap_input: bool = Form(...),
        driving_multiplier: float = Form(...),
        flag_stitching: bool = Form(...),
        flag_crop_driving_video_input: bool = Form(...),
        flag_video_editing_head_rotation: bool = Form(...),
        scale: float = Form(...),
        vx_ratio: float = Form(...),
        vy_ratio: float = Form(...),
        scale_crop_driving_video: float = Form(...),
        vx_ratio_crop_driving_video: float = Form(...),
        vy_ratio_crop_driving_video: float = Form(...),
        driving_smooth_observation_variance: float = Form(...)
):
    # 根据传入的表单参数构建 infer_params
    infer_params = LivePortraitParams(
        flag_is_animal=flag_is_animal,
        flag_pickle=flag_pickle,
        flag_relative_input=flag_relative_input,
        flag_do_crop_input=flag_do_crop_input,
        flag_remap_input=flag_remap_input,
        driving_multiplier=driving_multiplier,
        flag_stitching=flag_stitching,
        flag_crop_driving_video_input=flag_crop_driving_video_input,
        flag_video_editing_head_rotation=flag_video_editing_head_rotation,
        scale=scale,
        vx_ratio=vx_ratio,
        vy_ratio=vy_ratio,
        scale_crop_driving_video=scale_crop_driving_video,
        vx_ratio_crop_driving_video=vx_ratio_crop_driving_video,
        vy_ratio_crop_driving_video=vy_ratio_crop_driving_video,
        driving_smooth_observation_variance=driving_smooth_observation_variance
    )

    global pipe
    pipe.init_vars()
    if infer_params.flag_is_animal != pipe.is_animal:
        pipe.init_models(is_animal=infer_params.flag_is_animal)

    args_user = {
        'flag_relative_motion': infer_params.flag_relative_input,
        'flag_do_crop': infer_params.flag_do_crop_input,
        'flag_pasteback': infer_params.flag_remap_input,
        'driving_multiplier': infer_params.driving_multiplier,
        'flag_stitching': infer_params.flag_stitching,
        'flag_crop_driving_video': infer_params.flag_crop_driving_video_input,
        'flag_video_editing_head_rotation': infer_params.flag_video_editing_head_rotation,
        'src_scale': infer_params.scale,
        'src_vx_ratio': infer_params.vx_ratio,
        'src_vy_ratio': infer_params.vy_ratio,
        'dri_scale': infer_params.scale_crop_driving_video,
        'dri_vx_ratio': infer_params.vx_ratio_crop_driving_video,
        'dri_vy_ratio': infer_params.vy_ratio_crop_driving_video,
    }
    # update config from user input
    update_ret = pipe.update_cfg(args_user)

    # 保存 source_image 到指定目录
    temp_dir = os.path.join(result_dir, f"temp-{datetime.datetime.now().strftime('%Y-%m-%d-%H%M%S')}")
    os.makedirs(temp_dir, exist_ok=True)
    if source_image and source_image.filename:
        source_image_path = os.path.join(temp_dir, source_image.filename)
        with open(source_image_path, "wb") as buffer:
            buffer.write(await source_image.read())  # 将内容写入文件
    else:
        source_image_path = None

    if driving_video and driving_video.filename:
        driving_video_path = os.path.join(temp_dir, driving_video.filename)
        with open(driving_video_path, "wb") as buffer:
            buffer.write(await driving_video.read())  # 将内容写入文件
    else:
        driving_video_path = None

    if driving_pickle and driving_pickle.filename:
        driving_pickle_path = os.path.join(temp_dir, driving_pickle.filename)
        with open(driving_pickle_path, "wb") as buffer:
            buffer.write(await driving_pickle.read())  # 将内容写入文件
    else:
        driving_pickle_path = None

    save_dir = os.path.join(result_dir, f"{datetime.datetime.now().strftime('%Y-%m-%d-%H%M%S')}")
    os.makedirs(save_dir, exist_ok=True)

    if infer_params.flag_pickle:
        if source_image_path and driving_pickle_path:
            run_with_pkl(source_image_path, driving_pickle_path, save_dir)
    else:
        if source_image_path and driving_video_path:
            run_with_video(source_image_path, driving_video_path, save_dir)
    # zip all files and return
    # 使用 BytesIO 在内存中创建一个字节流
    zip_buffer = io.BytesIO()

    # 使用 ZipFile 将文件夹内容压缩到 zip_buffer 中
    with ZipFile(zip_buffer, "w") as zip_file:
        for root, dirs, files in os.walk(save_dir):
            for file in files:
                file_path = os.path.join(root, file)
                # 添加文件到 ZIP 文件中
                zip_file.write(file_path, arcname=os.path.relpath(file_path, save_dir))

    # 确保缓冲区指针在开始位置，以便读取整个内容
    zip_buffer.seek(0)
    shutil.rmtree(temp_dir)
    shutil.rmtree(save_dir)
    # 通过 StreamingResponse 返回 zip 文件
    return StreamingResponse(zip_buffer, media_type="application/zip",
                             headers={"Content-Disposition": "attachment; filename=output.zip"})


if __name__ == "__main__":
    import uvicorn

    uvicorn.run(app, host=os.environ.get("FLIP_IP", "127.0.0.1"), port=os.environ.get("FLIP_PORT", 9871))