fix: retargeting feature leakage

app.py

@@ -72,7 +72,7 @@ data_examples = [
 # Define components first
 eye_retargeting_slider = gr.Slider(minimum=0, maximum=0.8, step=0.01, label="target eyes-open ratio")
 lip_retargeting_slider = gr.Slider(minimum=0, maximum=0.8, step=0.01, label="target lip-open ratio")
-retargeting_input_image = gr.Image(type="numpy")
+retargeting_input_image = gr.Image(type="filepath")
 output_image = gr.Image(type="numpy")
 output_image_paste_back = gr.Image(type="numpy")
 output_video = gr.Video()
@@ -144,11 +144,11 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
             examples_per_page=5,
             cache_examples=False,
         )
-    gr.Markdown(load_description("assets/gradio_description_retargeting.md"), visible=False)
-    with gr.Row(visible=False):
+    gr.Markdown(load_description("assets/gradio_description_retargeting.md"), visible=True)
+    with gr.Row(visible=True):
         eye_retargeting_slider.render()
         lip_retargeting_slider.render()
-    with gr.Row(visible=False):
+    with gr.Row(visible=True):
         process_button_retargeting = gr.Button("🚗 Retargeting", variant="primary")
         process_button_reset_retargeting = gr.ClearButton(
             [
@@ -160,10 +160,21 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
             ],
             value="🧹 Clear"
         )
-    with gr.Row(visible=False):
+    with gr.Row(visible=True):
         with gr.Column():
             with gr.Accordion(open=True, label="Retargeting Input"):
                 retargeting_input_image.render()
+                gr.Examples(
+                    examples=[
+                        [osp.join(example_portrait_dir, "s9.jpg")],
+                        [osp.join(example_portrait_dir, "s6.jpg")],
+                        [osp.join(example_portrait_dir, "s10.jpg")],
+                        [osp.join(example_portrait_dir, "s5.jpg")],
+                        [osp.join(example_portrait_dir, "s7.jpg")],
+                    ],
+                    inputs=[retargeting_input_image],
+                    cache_examples=False,
+                )
         with gr.Column():
             with gr.Accordion(open=True, label="Retargeting Result"):
                 output_image.render()
@@ -174,7 +185,7 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
     process_button_retargeting.click(
         # fn=gradio_pipeline.execute_image,
         fn=gpu_wrapped_execute_image,
-        inputs=[eye_retargeting_slider, lip_retargeting_slider],
+        inputs=[eye_retargeting_slider, lip_retargeting_slider, retargeting_input_image, flag_do_crop_input],
         outputs=[output_image, output_image_paste_back],
         show_progress=True
     )
@@ -190,11 +201,11 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
         outputs=[output_video, output_video_concat],
         show_progress=True
     )
-    image_input.change(
-        fn=gradio_pipeline.prepare_retargeting,
-        inputs=image_input,
-        outputs=[eye_retargeting_slider, lip_retargeting_slider, retargeting_input_image]
-    )
+    # image_input.change(
+    #     fn=gradio_pipeline.prepare_retargeting,
+    #     inputs=image_input,
+    #     outputs=[eye_retargeting_slider, lip_retargeting_slider, retargeting_input_image]
+    # )
     video_input.upload(
         fn=is_square_video,
         inputs=video_input,

assets/gradio_description_retargeting.md

@@ -1 +1 @@
-<span style="font-size: 1.2em;">🔥 To change the target eyes-open and lip-open ratio of the source portrait, please drag the sliders and then click the <strong>🚗 Retargeting</strong> button. The result would be shown in the middle block. You can try running it multiple times. <strong>😊 Set both ratios to 0.8 to see what's going on!</strong> </span>
+<span style="font-size: 1.2em;">🔥 To change the eyes and lip open ratio of the source portrait, please drag the sliders and then click the <strong>🚗 Retargeting</strong> button. The result would be shown in the blocks. You can try running it multiple times. <strong>😊 Set both ratios to 0.8 to see what's going on!</strong> </span>

src/gradio_pipeline.py

@@ -26,16 +26,6 @@ class GradioPipeline(LivePortraitPipeline):
         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
-        self.x_s_user = None
-        self.source_lmk_user = None
-        self.mask_ori = None
-        self.img_rgb = None
-        self.crop_M_c2o = None
-
 
     def execute_video(
         self,
@@ -66,30 +56,23 @@ class GradioPipeline(LivePortraitPipeline):
         else:
             raise gr.Error("The input source portrait or driving video hasn't been prepared yet 💥!", duration=5)
 
-    def execute_image(self, input_eye_ratio: float, input_lip_ratio: float):
+    def execute_image(self, input_eye_ratio: float, input_lip_ratio: float, input_image, flag_do_crop = True):
         """ for single image retargeting
         """
+        # disposable feature
+        f_s_user, x_s_user, source_lmk_user, crop_M_c2o, mask_ori, img_rgb = \
+        self.prepare_retargeting(input_image, flag_do_crop)
+
         if input_eye_ratio is None or input_eye_ratio is None:
             raise gr.Error("Invalid ratio input 💥!", duration=5)
-        elif 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
-                )
         else:
-            x_s_user = self.x_s_user.to("cuda")
-            f_s_user = self.f_s_user.to("cuda")
+            x_s_user = x_s_user.to("cuda")
+            f_s_user = 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)
+            combined_eye_ratio_tensor = self.live_portrait_wrapper.calc_combined_eye_ratio([[input_eye_ratio]], 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)
+            combined_lip_ratio_tensor = self.live_portrait_wrapper.calc_combined_lip_ratio([[input_lip_ratio]], 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
@@ -97,21 +80,20 @@ class GradioPipeline(LivePortraitPipeline):
             # 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)
+            out_to_ori_blend = paste_back(out, crop_M_c2o, img_rgb, 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):
+    def prepare_retargeting(self, input_image, flag_do_crop = True):
         """ for single image retargeting
         """
-        if input_image_path is not None:
+        if input_image 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)
-            log(f"Load source image from {input_image_path}.")
+            img_rgb = load_img_online(input_image, mode='rgb', max_dim=1280, n=16)
+            log(f"Load source image from {input_image}.")
             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'])
@@ -120,23 +102,13 @@ class GradioPipeline(LivePortraitPipeline):
             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
-            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]))
-            # update slider
-            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]
-            return eye_close_ratio, lip_close_ratio, self.I_s_vis
+            f_s_user = self.live_portrait_wrapper.extract_feature_3d(I_s)
+            x_s_user = self.live_portrait_wrapper.transform_keypoint(x_s_info)
+            source_lmk_user = crop_info['lmk_crop']
+            crop_M_c2o = crop_info['M_c2o']
+            mask_ori = prepare_paste_back(inference_cfg.mask_crop, crop_info['M_c2o'], dsize=(img_rgb.shape[1], img_rgb.shape[0]))
+            return f_s_user, x_s_user, source_lmk_user, crop_M_c2o, mask_ori, img_rgb
         else:
             # when press the clear button, go here
-            return 0.8, 0.8, self.I_s_vis
+            raise gr.Error("The retargeting input hasn't been prepared yet 💥!", duration=5)
+