Update live_portrait_pipeline.py

liveportrait/live_portrait_pipeline.py

@@ -79,9 +79,10 @@ class LivePortraitPipeline(object):
                 source_np, i, total_frames, mismatch_method
             )
             driving_frame = driving_images_np[i]
-            
-            source_lmk = crop_info[i]["lmk_crop"]
-            img_crop_256x256 = crop_info[i]["img_crop_256x256"]
+
+            safe_index = min(i, len(crop_info) - 1)
+            source_lmk = crop_info[safe_index]["lmk_crop"]
+            img_crop_256x256 = crop_info[safe_index]["img_crop_256x256"]
 
             if inference_cfg.flag_do_crop:
                 I_s = self.live_portrait_wrapper.prepare_source(img_crop_256x256)
@@ -252,7 +253,7 @@ class LivePortraitPipeline(object):
             # Transform and blend
             I_p_i_to_ori = _transform_img(
                 I_p_i,
-                crop_info[i]["M_c2o"],
+                crop_info[safe_index]["M_c2o"],
                 dsize=(source_frame_rgb.shape[1], source_frame_rgb.shape[0]),
             )
 
@@ -261,7 +262,7 @@ class LivePortraitPipeline(object):
                     inference_cfg.mask_crop = cv2.imread(os.path.join(script_directory, "utils", "resources", "mask_template.png"), cv2.IMREAD_COLOR)
                 mask_ori = _transform_img(
                     inference_cfg.mask_crop,
-                    crop_info[i]["M_c2o"],
+                    crop_info[safe_index]["M_c2o"],
                     dsize=(source_frame_rgb.shape[1], source_frame_rgb.shape[0]),
                 )
                 mask_ori = mask_ori.astype(np.float32) / 255.0