src/gradio_pipeline.py

# coding: utf-8

"""
Pipeline for gradio

已适配 Hugging Face ZeroGPU：
1. 给会使用 CUDA 的 Gradio 回调函数添加 @spaces.GPU 装饰器。
2. 保留 Gradio 4.x 兼容写法。
3. 如果本地环境没有 spaces 包，也不会影响本地普通运行。
"""

import gradio as gr

# Hugging Face ZeroGPU 必须使用 spaces.GPU 装饰需要 GPU 的函数。
# 为了兼容本地运行，如果没有安装 spaces，则使用一个空装饰器。
try:
    import spaces
except Exception:
    class _DummySpaces:
        @staticmethod
        def GPU(duration=120):
            def decorator(func):
                return func
            return decorator

    spaces = _DummySpaces()

from .config.argument_config import ArgumentConfig
from .live_portrait_pipeline import LivePortraitPipeline
from .utils.io import load_img_online
from .utils.rprint import rlog as log
from .utils.crop import prepare_paste_back, paste_back
from .utils.camera import get_rotation_matrix
from .utils.retargeting_utils import calc_eye_close_ratio, calc_lip_close_ratio


def update_args(args, user_args):
    """
    根据用户输入更新参数。
    """
    for k, v in user_args.items():
        if hasattr(args, k):
            setattr(args, k, v)
    return args


class GradioPipeline(LivePortraitPipeline):

    def __init__(self, inference_cfg, crop_cfg, args: ArgumentConfig):
        super().__init__(inference_cfg, crop_cfg)

        self.args = args

        # 单图重定向状态缓存
        self.start_prepare = False
        self.f_s_user = None
        self.x_c_s_info_user = None
        self.x_s_user = None
        self.source_lmk_user = None
        self.mask_ori = None
        self.img_rgb = None
        self.crop_M_c2o = None
        self.I_s_vis = None

    @spaces.GPU(duration=300)
    def execute_video(
        self,
        input_image_path,
        input_video_path,
        flag_relative_input,
        flag_do_crop_input,
        flag_remap_input,
    ):
        """
        视频驱动肖像动画。

        注意：
        这个函数内部会调用模型推理和 CUDA，因此必须放在 @spaces.GPU 里。
        """
        if input_image_path is not None and input_video_path is not None:
            args_user = {
                "source_image": input_image_path,
                "driving_info": input_video_path,
                "flag_relative": flag_relative_input,
                "flag_do_crop": flag_do_crop_input,
                "flag_pasteback": flag_remap_input,
            }

            # 根据用户输入更新配置
            self.args = update_args(self.args, args_user)
            self.live_portrait_wrapper.update_config(self.args.__dict__)
            self.cropper.update_config(self.args.__dict__)

            # 执行视频驱动动画
            video_path, video_path_concat = self.execute(self.args)

            return video_path, video_path_concat

        raise gr.Error(
            "The input source portrait or driving video hasn't been prepared yet 💥!",
            duration=5,
        )

    @spaces.GPU(duration=180)
    def execute_image(self, input_eye_ratio: float, input_lip_ratio: float):
        """
        单图表情重定向。

        注意：
        这里会执行 .to("cuda")、retarget_eye、retarget_lip、warp_decode，
        因此必须放在 @spaces.GPU 里。
        """
        if input_eye_ratio is None or input_lip_ratio is None:
            raise gr.Error("Invalid ratio input 💥!", duration=5)

        if self.f_s_user is None:
            if self.start_prepare:
                raise gr.Error(
                    "The source portrait is under processing 💥! Please wait for a second.",
                    duration=5,
                )

            raise gr.Error(
                "The source portrait hasn't been prepared yet 💥! Please scroll to the top of the page to upload.",
                duration=5,
            )

        x_s_user = self.x_s_user.to("cuda")
        f_s_user = self.f_s_user.to("cuda")

        # 计算眼睛重定向
        combined_eye_ratio_tensor = self.live_portrait_wrapper.calc_combined_eye_ratio(
            [[input_eye_ratio]],
            self.source_lmk_user,
        )
        eyes_delta = self.live_portrait_wrapper.retarget_eye(
            x_s_user,
            combined_eye_ratio_tensor,
        )

        # 计算嘴唇重定向
        combined_lip_ratio_tensor = self.live_portrait_wrapper.calc_combined_lip_ratio(
            [[input_lip_ratio]],
            self.source_lmk_user,
        )
        lip_delta = self.live_portrait_wrapper.retarget_lip(
            x_s_user,
            combined_lip_ratio_tensor,
        )

        num_kp = x_s_user.shape[1]

        # 默认基于 x_s 做变形
        x_d_new = (
            x_s_user
            + eyes_delta.reshape(-1, num_kp, 3)
            + lip_delta.reshape(-1, num_kp, 3)
        )

        # 解码输出
        out = self.live_portrait_wrapper.warp_decode(
            f_s_user,
            x_s_user,
            x_d_new,
        )
        out = self.live_portrait_wrapper.parse_output(out["out"])[0]

        # 贴回原图
        out_to_ori_blend = paste_back(
            out,
            self.crop_M_c2o,
            self.img_rgb,
            self.mask_ori,
        )

        return out, out_to_ori_blend

    @spaces.GPU(duration=180)
    def prepare_retargeting(self, input_image_path, flag_do_crop=True):
        """
        单图表情重定向的预处理。

        注意：
        日志中的报错发生在这个函数调用链里：
        prepare_retargeting -> prepare_source -> x.cuda(...)
        所以这个函数必须使用 @spaces.GPU。
        """
        if input_image_path is not None:
            self.start_prepare = True

            inference_cfg = self.live_portrait_wrapper.cfg

            # 读取源图
            img_rgb = load_img_online(
                input_image_path,
                mode="rgb",
                max_dim=1280,
                n=16,
            )
            log(f"Load source image from {input_image_path}.")

            # 裁剪人脸
            crop_info = self.cropper.crop_single_image(img_rgb)

            if flag_do_crop:
                I_s = self.live_portrait_wrapper.prepare_source(
                    crop_info["img_crop_256x256"]
                )
            else:
                I_s = self.live_portrait_wrapper.prepare_source(img_rgb)

            # 提取关键点信息
            x_s_info = self.live_portrait_wrapper.get_kp_info(I_s)

            # 保留原逻辑：计算旋转矩阵
            # 当前变量暂未在后续使用，但保留，避免影响原项目行为。
            _ = get_rotation_matrix(
                x_s_info["pitch"],
                x_s_info["yaw"],
                x_s_info["roll"],
            )

            # 缓存后续单图重定向需要的数据
            self.f_s_user = self.live_portrait_wrapper.extract_feature_3d(I_s)
            self.x_s_user = self.live_portrait_wrapper.transform_keypoint(x_s_info)
            self.x_s_info_user = x_s_info
            self.source_lmk_user = crop_info["lmk_crop"]
            self.img_rgb = img_rgb
            self.crop_M_c2o = crop_info["M_c2o"]
            self.mask_ori = prepare_paste_back(
                inference_cfg.mask_crop,
                crop_info["M_c2o"],
                dsize=(img_rgb.shape[1], img_rgb.shape[0]),
            )

            # 更新滑块默认值
            eye_close_ratio = calc_eye_close_ratio(self.source_lmk_user[None])
            eye_close_ratio = float(eye_close_ratio.squeeze(0).mean())

            lip_close_ratio = calc_lip_close_ratio(self.source_lmk_user[None])
            lip_close_ratio = float(lip_close_ratio.squeeze(0).mean())

            # 预览图
            self.I_s_vis = self.live_portrait_wrapper.parse_output(I_s)[0]

            self.start_prepare = False

            return eye_close_ratio, lip_close_ratio, self.I_s_vis

        # 点击清空按钮时走这里
        if self.I_s_vis is not None:
            return 0.8, 0.8, self.I_s_vis

        return 0.8, 0.8, None