Update src/gradio_pipeline.py

src/gradio_pipeline.py

@@ -2,8 +2,29 @@
 
 """
 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
@@ -12,21 +33,25 @@ 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):
-    """update the args according to user inputs
+    """
+    根据用户输入更新参数。
     """
     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.live_portrait_wrapper = self.live_portrait_wrapper
+
         self.args = args
-        # for single image retargeting
+
+        # 单图重定向状态缓存
         self.start_prepare = False
         self.f_s_user = None
         self.x_c_s_info_user = None
@@ -35,8 +60,9 @@ class GradioPipeline(LivePortraitPipeline):
         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,
@@ -44,99 +70,185 @@ class GradioPipeline(LivePortraitPipeline):
         flag_relative_input,
         flag_do_crop_input,
         flag_remap_input,
-        ):
-        """ for video driven potrait animation
+    ):
+        """
+        视频驱动肖像动画。
+
+        注意：
+        这个函数内部会调用模型推理和 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,
+                "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,
             }
-            # update config from user 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 driven animation
+
+            # 执行视频驱动动画
             video_path, video_path_concat = self.execute(self.args)
-            # gr.Info("Run successfully!", duration=2)
-            return video_path, video_path_concat,
-        else:
-            raise gr.Error("The input source portrait or driving video hasn't been prepared yet 💥!", duration=5)
 
+            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):
-        """ for single image retargeting
         """
-        if input_eye_ratio is None or input_eye_ratio is None:
+        单图表情重定向。
+
+        注意：
+        这里会执行 .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)
-        elif self.f_s_user is None:
+
+        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
-                )
-            else:
-                raise gr.Error(
-                    "The source portrait hasn't been prepared yet 💥! Please scroll to the top of the page to upload.",
-                    duration=5
+                    duration=5,
                 )
-        else:
-            x_s_user = self.x_s_user.to("cuda")
-            f_s_user = self.f_s_user.to("cuda")
-            # ∆_eyes,i = R_eyes(x_s; c_s,eyes, c_d,eyes,i)
-            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)
-            # ∆_lip,i = R_lip(x_s; c_s,lip, c_d,lip,i)
-            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]
-            # default: use x_s
-            x_d_new = x_s_user + eyes_delta.reshape(-1, num_kp, 3) + lip_delta.reshape(-1, num_kp, 3)
-            # D(W(f_s; x_s, x′_d))
-            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)
-            # gr.Info("Run successfully!", duration=2)
-            return out, out_to_ori_blend
-
-
-    def prepare_retargeting(self, input_image_path, flag_do_crop = True):
-        """ for single image retargeting
+
+            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:
-            # gr.Info("Upload successfully!", duration=2)
             self.start_prepare = True
+
             inference_cfg = self.live_portrait_wrapper.cfg
-            ######## process source portrait ########
-            img_rgb = load_img_online(input_image_path, mode='rgb', max_dim=1280, n=16)
+
+            # 读取源图
+            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'])
+                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)
-            R_s = get_rotation_matrix(x_s_info['pitch'], x_s_info['yaw'], x_s_info['roll'])
-            ############################################
 
-            # record global info for next time use
+            # 保留原逻辑：计算旋转矩阵
+            # 当前变量暂未在后续使用，但保留，避免影响原项目行为。
+            _ = 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.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]))
-            # update slider
+            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())
-            # for vis
+
+            # 预览图
             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
-        else:
-            # when press the clear button, go here
+
+        # 点击清空按钮时走这里
+        if self.I_s_vis is not None:
             return 0.8, 0.8, self.I_s_vis
+
+        return 0.8, 0.8, None