avoid vram issues

wan/modules/animate/preprocess/process_pipepline.py

@@ -1,585 +1,585 @@
-# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
-import os
-import numpy as np
-import shutil
-import torch
-from diffusers import FluxKontextPipeline
-import cv2
-from loguru import logger
-from PIL import Image
-try:
-    import moviepy.editor as mpy
-except:
-    import moviepy as mpy
-
-from decord import VideoReader
-from pose2d import Pose2d
-from pose2d_utils import AAPoseMeta
-from utils import resize_by_area, get_frame_indices, padding_resize, get_face_bboxes, get_aug_mask, get_mask_body_img
-from human_visualization import draw_aapose_by_meta_new
-from retarget_pose import get_retarget_pose
-import sam2.modeling.sam.transformer as transformer
-transformer.USE_FLASH_ATTN = True
-transformer.MATH_KERNEL_ON = False
-transformer.OLD_GPU = False
-from sam_utils import build_sam2_video_predictor
-
-
-class ProcessPipeline():
-    def __init__(self, det_checkpoint_path, pose2d_checkpoint_path, sam_checkpoint_path, flux_kontext_path):
-        self.pose2d = Pose2d(checkpoint=pose2d_checkpoint_path, detector_checkpoint=det_checkpoint_path)
-
-        model_cfg = "sam2_hiera_l.yaml"
-        if sam_checkpoint_path is not None:
-            self.predictor = build_sam2_video_predictor(model_cfg, sam_checkpoint_path)    
-        if flux_kontext_path is not None:
-            self.flux_kontext = FluxKontextPipeline.from_pretrained(flux_kontext_path, torch_dtype=torch.bfloat16).to("cuda")
-
-    def __call__(self, video_path, refer_image_path, output_path, resolution_area=[1280, 720], fps=30, iterations=3, k=7, w_len=1, h_len=1, retarget_flag=False, use_flux=False, replace_flag=False):
-        if replace_flag:
-
-            video_reader = VideoReader(video_path)
-            frame_num = len(video_reader)
-            print('frame_num: {}'.format(frame_num))
-            
-            video_fps = video_reader.get_avg_fps()
-            print('video_fps: {}'.format(video_fps))
-            print('fps: {}'.format(fps))
-
-            # TODO: Maybe we can switch to PyAV later, which can get accurate frame num
-            duration = video_reader.get_frame_timestamp(-1)[-1]      
-            expected_frame_num = int(duration * video_fps + 0.5) 
-            ratio = abs((frame_num - expected_frame_num)/frame_num)         
-            if ratio > 0.1:
-                print("Warning: The difference between the actual number of frames and the expected number of frames is two large")
-                frame_num = expected_frame_num
-
-            if fps == -1:
-                fps = video_fps
-
-            target_num = int(frame_num / video_fps * fps)
-            print('target_num: {}'.format(target_num))
-            idxs = get_frame_indices(frame_num, video_fps, target_num, fps)
-            frames = video_reader.get_batch(idxs).asnumpy()
-
-            frames = [resize_by_area(frame, resolution_area[0] * resolution_area[1], divisor=16) for frame in frames]
-            height, width = frames[0].shape[:2]
-            logger.info(f"Processing pose meta")
-
-
-            tpl_pose_metas = self.pose2d(frames)
-
-            face_images = []
-            for idx, meta in enumerate(tpl_pose_metas):
-                face_bbox_for_image = get_face_bboxes(meta['keypoints_face'][:, :2], scale=1.3,
-                                                    image_shape=(frames[0].shape[0], frames[0].shape[1]))
-
-                x1, x2, y1, y2 = face_bbox_for_image
-                face_image = frames[idx][y1:y2, x1:x2]
-                face_image = cv2.resize(face_image, (512, 512))
-                face_images.append(face_image)
-
-            logger.info(f"Processing reference image: {refer_image_path}")
-            refer_img = cv2.imread(refer_image_path)
-            src_ref_path = os.path.join(output_path, 'src_ref.png')
-            shutil.copy(refer_image_path, src_ref_path)
-            refer_img = refer_img[..., ::-1]
-
-            refer_img = padding_resize(refer_img, height, width)
-            logger.info(f"Processing template video: {video_path}")
-            tpl_retarget_pose_metas = [AAPoseMeta.from_humanapi_meta(meta) for meta in tpl_pose_metas]
-            cond_images = []
-
-            for idx, meta in enumerate(tpl_retarget_pose_metas):
-                canvas = np.zeros_like(refer_img)
-                conditioning_image = draw_aapose_by_meta_new(canvas, meta)
-                cond_images.append(conditioning_image)
-            masks = self.get_mask_from_face_bbox(frames, 400, tpl_pose_metas)
-
-            bg_images = []
-            aug_masks = []
-
-            for frame, mask in zip(frames, masks):
-                if iterations > 0:
-                    _, each_mask = get_mask_body_img(frame, mask, iterations=iterations, k=k)
-                    each_aug_mask = get_aug_mask(each_mask, w_len=w_len, h_len=h_len)
-                else:
-                    each_aug_mask = mask
-
-                each_bg_image = frame * (1 - each_aug_mask[:, :, None])
-                bg_images.append(each_bg_image)
-                aug_masks.append(each_aug_mask)
-
-            src_face_path = os.path.join(output_path, 'src_face.mp4')
-            mpy.ImageSequenceClip(face_images, fps=fps).write_videofile(src_face_path)
-
-            src_pose_path = os.path.join(output_path, 'src_pose.mp4')
-            mpy.ImageSequenceClip(cond_images, fps=fps).write_videofile(src_pose_path)
-
-            src_bg_path = os.path.join(output_path, 'src_bg.mp4')
-            mpy.ImageSequenceClip(bg_images, fps=fps).write_videofile(src_bg_path)
-
-            aug_masks_new = [np.stack([mask * 255, mask * 255, mask * 255], axis=2) for mask in aug_masks]
-            src_mask_path = os.path.join(output_path, 'src_mask.mp4')
-            mpy.ImageSequenceClip(aug_masks_new, fps=fps).write_videofile(src_mask_path)
-            return True
-        else:
-            logger.info(f"Processing reference image: {refer_image_path}")
-            refer_img = cv2.imread(refer_image_path)
-            src_ref_path = os.path.join(output_path, 'src_ref.png')
-            shutil.copy(refer_image_path, src_ref_path)
-            refer_img = refer_img[..., ::-1]
-            
-            refer_img = resize_by_area(refer_img, resolution_area[0] * resolution_area[1], divisor=16)
-            
-            refer_pose_meta = self.pose2d([refer_img])[0]
-
-
-            logger.info(f"Processing template video: {video_path}")
-            video_reader = VideoReader(video_path)
-            frame_num = len(video_reader)
-            print('frame_num: {}'.format(frame_num))
-
-            video_fps = video_reader.get_avg_fps()
-            print('video_fps: {}'.format(video_fps))
-            print('fps: {}'.format(fps))
-
-            # TODO: Maybe we can switch to PyAV later, which can get accurate frame num
-            duration = video_reader.get_frame_timestamp(-1)[-1]      
-            expected_frame_num = int(duration * video_fps + 0.5) 
-            ratio = abs((frame_num - expected_frame_num)/frame_num)         
-            if ratio > 0.1:
-                print("Warning: The difference between the actual number of frames and the expected number of frames is two large")
-                frame_num = expected_frame_num
-
-            if fps == -1:
-                fps = video_fps
-                
-            target_num = int(frame_num / video_fps * fps)
-            print('target_num: {}'.format(target_num))
-            idxs = get_frame_indices(frame_num, video_fps, target_num, fps)
-            frames = video_reader.get_batch(idxs).asnumpy()
-
-            logger.info(f"Processing pose meta")
-
-            tpl_pose_meta0 = self.pose2d(frames[:1])[0]
-            tpl_pose_metas = self.pose2d(frames)
-
-            face_images = []
-            for idx, meta in enumerate(tpl_pose_metas):
-                face_bbox_for_image = get_face_bboxes(meta['keypoints_face'][:, :2], scale=1.3,
-                                                    image_shape=(frames[0].shape[0], frames[0].shape[1]))
-
-                x1, x2, y1, y2 = face_bbox_for_image
-                face_image = frames[idx][y1:y2, x1:x2]
-                face_image = cv2.resize(face_image, (512, 512))
-                face_images.append(face_image)
-
-            if retarget_flag:
-                if use_flux:
-                    tpl_prompt, refer_prompt = self.get_editing_prompts(tpl_pose_metas, refer_pose_meta)
-                    refer_input = Image.fromarray(refer_img)
-                    refer_edit = self.flux_kontext(
-                            image=refer_input,
-                            height=refer_img.shape[0],
-                            width=refer_img.shape[1],
-                            prompt=refer_prompt,
-                            guidance_scale=2.5,
-                            num_inference_steps=28,
-                        ).images[0]
-                    
-                    refer_edit = Image.fromarray(padding_resize(np.array(refer_edit), refer_img.shape[0], refer_img.shape[1]))
-                    refer_edit_path = os.path.join(output_path, 'refer_edit.png')
-                    refer_edit.save(refer_edit_path)
-                    refer_edit_pose_meta = self.pose2d([np.array(refer_edit)])[0]
-
-                    tpl_img = frames[1]
-                    tpl_input = Image.fromarray(tpl_img)
-                    
-                    tpl_edit = self.flux_kontext(
-                            image=tpl_input,
-                            height=tpl_img.shape[0],
-                            width=tpl_img.shape[1],
-                            prompt=tpl_prompt,
-                            guidance_scale=2.5,
-                            num_inference_steps=28,
-                        ).images[0]
-                    
-                    tpl_edit = Image.fromarray(padding_resize(np.array(tpl_edit), tpl_img.shape[0], tpl_img.shape[1]))
-                    tpl_edit_path = os.path.join(output_path, 'tpl_edit.png')
-                    tpl_edit.save(tpl_edit_path)
-                    tpl_edit_pose_meta0 = self.pose2d([np.array(tpl_edit)])[0]
-                    tpl_retarget_pose_metas = get_retarget_pose(tpl_pose_meta0, refer_pose_meta, tpl_pose_metas, tpl_edit_pose_meta0, refer_edit_pose_meta)
-                else:
-                    tpl_retarget_pose_metas = get_retarget_pose(tpl_pose_meta0, refer_pose_meta, tpl_pose_metas, None, None)
-            else:
-               tpl_retarget_pose_metas = [AAPoseMeta.from_humanapi_meta(meta) for meta in tpl_pose_metas]
-
-            cond_images = []
-            for idx, meta in enumerate(tpl_retarget_pose_metas):
-                if retarget_flag:
-                    canvas = np.zeros_like(refer_img)
-                    conditioning_image = draw_aapose_by_meta_new(canvas, meta)
-                else:
-                    canvas = np.zeros_like(frames[0])
-                    conditioning_image = draw_aapose_by_meta_new(canvas, meta)
-                    conditioning_image = padding_resize(conditioning_image, refer_img.shape[0], refer_img.shape[1])
-
-                cond_images.append(conditioning_image)
-
-            src_face_path = os.path.join(output_path, 'src_face.mp4')
-            mpy.ImageSequenceClip(face_images, fps=fps).write_videofile(src_face_path)
-
-            src_pose_path = os.path.join(output_path, 'src_pose.mp4')
-            mpy.ImageSequenceClip(cond_images, fps=fps).write_videofile(src_pose_path)
-            return True
-
-    def get_editing_prompts(self, tpl_pose_metas, refer_pose_meta):
-        arm_visible = False
-        leg_visible = False
-        for tpl_pose_meta in tpl_pose_metas:
-            tpl_keypoints = tpl_pose_meta['keypoints_body']
-            if tpl_keypoints[3].all() != 0 or tpl_keypoints[4].all() != 0 or tpl_keypoints[6].all() != 0 or tpl_keypoints[7].all() != 0:
-                if (tpl_keypoints[3][0] <= 1 and tpl_keypoints[3][1] <= 1 and tpl_keypoints[3][2] >= 0.75) or (tpl_keypoints[4][0] <= 1 and tpl_keypoints[4][1] <= 1 and tpl_keypoints[4][2] >= 0.75) or \
-                    (tpl_keypoints[6][0] <= 1 and tpl_keypoints[6][1] <= 1 and tpl_keypoints[6][2] >= 0.75) or (tpl_keypoints[7][0] <= 1 and tpl_keypoints[7][1] <= 1 and tpl_keypoints[7][2] >= 0.75):
-                    arm_visible = True
-            if tpl_keypoints[9].all() != 0 or tpl_keypoints[12].all() != 0 or tpl_keypoints[10].all() != 0 or tpl_keypoints[13].all() != 0:
-                if (tpl_keypoints[9][0] <= 1 and tpl_keypoints[9][1] <= 1 and tpl_keypoints[9][2] >= 0.75) or (tpl_keypoints[12][0] <= 1 and tpl_keypoints[12][1] <= 1 and tpl_keypoints[12][2] >= 0.75) or \
-                    (tpl_keypoints[10][0] <= 1 and tpl_keypoints[10][1] <= 1 and tpl_keypoints[10][2] >= 0.75) or (tpl_keypoints[13][0] <= 1 and tpl_keypoints[13][1] <= 1 and tpl_keypoints[13][2] >= 0.75):
-                    leg_visible = True
-            if arm_visible and leg_visible:
-                break
-        
-        if leg_visible:
-            if tpl_pose_meta['width'] > tpl_pose_meta['height']:
-                tpl_prompt = "Change the person to a standard T-pose (facing forward with arms extended). The person is standing. Feet and Hands are visible in the image."
-            else:
-                tpl_prompt = "Change the person to a standard pose with the face oriented forward and arms extending straight down by the sides. The person is standing. Feet and Hands are visible in the image."
-
-            if refer_pose_meta['width'] > refer_pose_meta['height']:
-                refer_prompt = "Change the person to a standard T-pose (facing forward with arms extended). The person is standing. Feet and Hands are visible in the image."
-            else:
-                refer_prompt = "Change the person to a standard pose with the face oriented forward and arms extending straight down by the sides. The person is standing. Feet and Hands are visible in the image."
-        elif arm_visible:
-            if tpl_pose_meta['width'] > tpl_pose_meta['height']:
-                tpl_prompt = "Change the person to a standard T-pose (facing forward with arms extended). Hands are visible in the image."
-            else:
-                tpl_prompt = "Change the person to a standard pose with the face oriented forward and arms extending straight down by the sides. Hands are visible in the image."
-
-            if refer_pose_meta['width'] > refer_pose_meta['height']:
-                refer_prompt = "Change the person to a standard T-pose (facing forward with arms extended). Hands are visible in the image."
-            else:
-                refer_prompt = "Change the person to a standard pose with the face oriented forward and arms extending straight down by the sides. Hands are visible in the image."
-        else:
-            tpl_prompt = "Change the person to face forward."
-            refer_prompt = "Change the person to face forward."
-
-        return tpl_prompt, refer_prompt
-    
-
-    def get_mask(self, frames, th_step, kp2ds_all):
-        frame_num = len(frames)
-        if frame_num < th_step:
-            num_step = 1
-        else:
-            num_step = (frame_num + th_step) // th_step
-
-        all_mask = []
-        for index in range(num_step):
-            each_frames = frames[index * th_step:(index + 1) * th_step]
-    
-            kp2ds = kp2ds_all[index * th_step:(index + 1) * th_step]
-            if len(each_frames) > 4:
-                key_frame_num = 4
-            elif 4 >= len(each_frames) > 0:
-                key_frame_num = 1
-            else:
-                continue
-
-            key_frame_step = len(kp2ds) // key_frame_num
-            key_frame_index_list = list(range(0, len(kp2ds), key_frame_step))
-
-            key_points_index = [0, 1, 2, 5, 8, 11, 10, 13]
-            key_frame_body_points_list = []
-            for key_frame_index in key_frame_index_list:
-                keypoints_body_list = []
-                body_key_points = kp2ds[key_frame_index]['keypoints_face']
-                for each_index in key_points_index:
-                    each_keypoint = body_key_points[each_index]
-                    if None is each_keypoint:
-                        continue
-                    keypoints_body_list.append(each_keypoint)
-
-                keypoints_body = np.array(keypoints_body_list)[:, :2]
-                wh = np.array([[kp2ds[0]['width'], kp2ds[0]['height']]])
-                points = (keypoints_body * wh).astype(np.int32)
-                key_frame_body_points_list.append(points)
-
-            inference_state = self.predictor.init_state_v2(frames=each_frames)
-            self.predictor.reset_state(inference_state)
-            ann_obj_id = 1
-            for ann_frame_idx, points in zip(key_frame_index_list, key_frame_body_points_list):
-                labels = np.array([1] * points.shape[0], np.int32)
-                _, out_obj_ids, out_mask_logits = self.predictor.add_new_points(
-                    inference_state=inference_state,
-                    frame_idx=ann_frame_idx,
-                    obj_id=ann_obj_id,
-                    points=points,
-                    labels=labels,
-                )
-
-            video_segments = {}
-            for out_frame_idx, out_obj_ids, out_mask_logits in self.predictor.propagate_in_video(inference_state):
-                video_segments[out_frame_idx] = {
-                    out_obj_id: (out_mask_logits[i] > 0.0).cpu().numpy()
-                    for i, out_obj_id in enumerate(out_obj_ids)
-                }
-
-            for out_frame_idx in range(len(video_segments)):
-                for out_obj_id, out_mask in video_segments[out_frame_idx].items():
-                    out_mask = out_mask[0].astype(np.uint8)
-                    all_mask.append(out_mask)
-
-        return all_mask
-    
-    def convert_list_to_array(self, metas):
-        metas_list = []
-        for meta in metas:
-            for key, value in meta.items():
-                if type(value) is list:
-                    value = np.array(value)
-                meta[key] = value
-            metas_list.append(meta)
-        return metas_list
-
-    def get_mask_from_face_bbox(self, frames, th_step, kp2ds_all):
-        """
-        Build masks using a face bounding box per key frame (derived from keypoints_face),
-        then propagate with SAM2 across each chunk of frames.
-        """
-        H, W = frames[0].shape[:2]
-
-        def _clip_box(x1, y1, x2, y2, W, H):
-            x1 = max(0, min(int(x1), W - 1))
-            x2 = max(0, min(int(x2), W - 1))
-            y1 = max(0, min(int(y1), H - 1))
-            y2 = max(0, min(int(y2), H - 1))
-            if x2 <= x1: x2 = min(W - 1, x1 + 1)
-            if y2 <= y1: y2 = min(H - 1, y1 + 1)
-            return x1, y1, x2, y2
-
-        frame_num = len(frames)
-        if frame_num < th_step:
-            num_step = 1
-        else:
-            num_step = (frame_num + th_step) // th_step
-
-        all_mask = []
-
-        for step_idx in range(num_step):
-            each_frames = frames[step_idx * th_step:(step_idx + 1) * th_step]
-            kp2ds = kp2ds_all[step_idx * th_step:(step_idx + 1) * th_step]
-            if len(each_frames) == 0:
-                continue
-
-            # pick a few key frames in this chunk
-            key_frame_num = 4 if len(each_frames) > 4 else 1
-            key_frame_step = max(1, len(kp2ds) // key_frame_num)
-            key_frame_index_list = list(range(0, len(kp2ds), key_frame_step))[:key_frame_num]
-
-            # compute face boxes on the selected key frames
-            key_frame_boxes = []
-            for kfi in key_frame_index_list:
-                meta = kp2ds[kfi]
-                # get_face_bboxes returns (x1, x2, y1, y2) in your code
-                x1, x2, y1, y2 = get_face_bboxes(
-                    meta['keypoints_face'][:, :2],
-                    scale=1.3,
-                    image_shape=(H, W)
-                )
-                x1, y1, x2, y2 = _clip_box(x1, y1, x2, y2, W, H)
-                key_frame_boxes.append(np.array([x1, y1, x2, y2], dtype=np.float32))
-
-            # init SAM2 for this chunk
-            with torch.no_grad(), torch.cuda.amp.autocast(dtype=torch.float16):
-                inference_state = self.predictor.init_state_v2(frames=each_frames)
-                self.predictor.reset_state(inference_state)
-                ann_obj_id = 1
-
-                # seed with box prompts (preferred), else fall back to points
-                for ann_frame_idx, box_xyxy in zip(key_frame_index_list, key_frame_boxes):
-                    used_box = False
-                    try:
-                        # If your predictor exposes a box API, this is ideal.
-                        _ = self.predictor.add_new_box(
-                            inference_state=inference_state,
-                            frame_idx=ann_frame_idx,
-                            obj_id=ann_obj_id,
-                            box=box_xyxy[None, :]  # shape (1, 4)
-                        )
-                        used_box = True
-                    except Exception:
-                        used_box = False
-
-                    if not used_box:
-                        # Fallback: sample a few positive points inside the box
-                        x1, y1, x2, y2 = box_xyxy.astype(int)
-                        cx, cy = (x1 + x2) // 2, (y1 + y2) // 2
-                        pts = np.array([
-                            [cx, cy],
-                            [x1 + (x2 - x1) // 4, cy],
-                            [x2 - (x2 - x1) // 4, cy],
-                            [cx, y1 + (y2 - y1) // 4],
-                            [cx, y2 - (y2 - y1) // 4],
-                        ], dtype=np.int32)
-                        labels = np.ones(len(pts), dtype=np.int32)  # 1 = positive
-                        _ = self.predictor.add_new_points(
-                            inference_state=inference_state,
-                            frame_idx=ann_frame_idx,
-                            obj_id=ann_obj_id,
-                            points=pts,
-                            labels=labels,
-                        )
-
-                # propagate across the chunk
-                video_segments = {}
-                for out_frame_idx, out_obj_ids, out_mask_logits in self.predictor.propagate_in_video(inference_state):
-                    video_segments[out_frame_idx] = {
-                        out_obj_id: (out_mask_logits[i] > 0.0).cpu().numpy()
-                        for i, out_obj_id in enumerate(out_obj_ids)
-                    }
-
-            # collect masks (single object id)
-            for out_frame_idx in range(len(video_segments)):
-                # (H, W) boolean/uint8
-                mask = next(iter(video_segments[out_frame_idx].values()))
-                mask = mask[0].astype(np.uint8)
-                all_mask.append(mask)
-
-        return all_mask
-    def get_mask_from_face_point(self, frames, th_step, kp2ds_all):
-        """
-        Build masks using a single face *center point* per key frame,
-        then propagate with SAM2 across each chunk of frames.
-        """
-        H, W = frames[0].shape[:2]
-
-        frame_num = len(frames)
-        num_step = 1 if frame_num < th_step else (frame_num + th_step) // th_step
-
-        all_mask = []
-
-        for step_idx in range(num_step):
-            each_frames = frames[step_idx * th_step:(step_idx + 1) * th_step]
-            kp2ds = kp2ds_all[step_idx * th_step:(step_idx + 1) * th_step]
-            if len(each_frames) == 0:
-                continue
-
-            # choose a few key frames to seed the object
-            key_frame_num = 1
-            key_frame_step = max(1, len(kp2ds) // key_frame_num)
-            key_frame_index_list = list(range(0, len(kp2ds), key_frame_step))[:key_frame_num]
-
-            # compute center point from face bbox for each selected key frame
-            center_pts = []
-            for kfi in key_frame_index_list:
-                meta = kp2ds[kfi]
-                # your helper returns (x1, x2, y1, y2)
-                x1, x2, y1, y2 = get_face_bboxes(
-                    meta['keypoints_face'][:, :2],
-                    scale=1.3,
-                    image_shape=(H, W)
-                )
-                cx = (x1 + x2) // 2
-                cy = (y1 + y2) // 2
-                # clip just in case
-                cx = int(max(0, min(cx, W - 1)))
-                cy = int(max(0, min(cy, H - 1)))
-                center_pts.append(np.array([cx, cy], dtype=np.int32))
-
-            # init SAM2 for this chunk
-            inference_state = self.predictor.init_state_v2(frames=each_frames)
-            self.predictor.reset_state(inference_state)
-            ann_obj_id = 1
-
-            # seed each key frame with a single positive point at the face center
-            for ann_frame_idx, pt in zip(key_frame_index_list, center_pts):
-                pts = pt[None, :]  # shape (1, 2)
-                labels = np.ones(1, dtype=np.int32)  # 1 = positive
-                _ = self.predictor.add_new_points(
-                    inference_state=inference_state,
-                    frame_idx=ann_frame_idx,
-                    obj_id=ann_obj_id,
-                    points=pts,
-                    labels=labels,
-                )
-
-            # propagate across the chunk
-            video_segments = {}
-            for out_frame_idx, out_obj_ids, out_mask_logits in self.predictor.propagate_in_video(inference_state):
-                video_segments[out_frame_idx] = {
-                    out_obj_id: (out_mask_logits[i] > 0.0).cpu().numpy()
-                    for i, out_obj_id in enumerate(out_obj_ids)
-                }
-
-            # collect masks (single object id)
-            for out_frame_idx in range(len(video_segments)):
-                mask = next(iter(video_segments[out_frame_idx].values()))
-                mask = mask[0].astype(np.uint8)
-                all_mask.append(mask)
-
-        return all_mask
-
-    def get_face_bbox_masks(self, frames, kp2ds_all, scale=1.3, feather_px=0, keep_soft=False):
-        """
-        Create a per-frame mask that's simply the face bounding box.
-        - scale: bbox scale factor used by get_face_bboxes
-        - feather_px: optional Gaussian blur in pixels to feather edges
-        - keep_soft: if True, keep float [0,1] soft mask (after blur); else binarize to {0,1}
-        """
-        H, W = frames[0].shape[:2]
-
-        def _clip_box(x1, y1, x2, y2):
-            x1 = max(0, min(int(x1), W - 1))
-            x2 = max(0, min(int(x2), W - 1))
-            y1 = max(0, min(int(y1), H - 1))
-            y2 = max(0, min(int(y2), H - 1))
-            if x2 <= x1: x2 = min(W - 1, x1 + 1)
-            if y2 <= y1: y2 = min(H - 1, y1 + 1)
-            return x1, y1, x2, y2
-
-        masks = []
-        last_box = None
-        for meta in kp2ds_all:
-            # get_face_bboxes returns (x1, x2, y1, y2)
-            try:
-                x1, x2, y1, y2 = get_face_bboxes(
-                    meta['keypoints_face'][:, :2],
-                    scale=scale,
-                    image_shape=(H, W)
-                )
-                x1, y1, x2, y2 = _clip_box(x1, y1, x2, y2)
-                last_box = (x1, y1, x2, y2)
-            except Exception:
-                # fallback: reuse last seen box to avoid holes
-                if last_box is None:
-                    # no detection yet: push empty mask
-                    masks.append(np.zeros((H, W), dtype=np.uint8))
-                    continue
-                x1, y1, x2, y2 = last_box
-
-            m = np.zeros((H, W), dtype=np.float32)
-            m[y1:y2, x1:x2] = 1.0
-
-            if feather_px and feather_px > 0:
-                # kernel size must be odd and >= 3
-                k = max(3, int(feather_px) | 1)
-                m = cv2.GaussianBlur(m, (k, k), 0)
-
-            if keep_soft:
-                masks.append(m)  # float [0,1]
-            else:
-                masks.append((m >= 0.5).astype(np.uint8))  # hard {0,1}
-
-        return masks
-
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+import os
+import numpy as np
+import shutil
+import torch
+from diffusers import FluxKontextPipeline
+import cv2
+from loguru import logger
+from PIL import Image
+try:
+    import moviepy.editor as mpy
+except:
+    import moviepy as mpy
+
+from decord import VideoReader
+from pose2d import Pose2d
+from pose2d_utils import AAPoseMeta
+from utils import resize_by_area, get_frame_indices, padding_resize, get_face_bboxes, get_aug_mask, get_mask_body_img
+from human_visualization import draw_aapose_by_meta_new
+from retarget_pose import get_retarget_pose
+import sam2.modeling.sam.transformer as transformer
+transformer.USE_FLASH_ATTN = True
+transformer.MATH_KERNEL_ON = True
+transformer.OLD_GPU = False
+from sam_utils import build_sam2_video_predictor
+
+
+class ProcessPipeline():
+    def __init__(self, det_checkpoint_path, pose2d_checkpoint_path, sam_checkpoint_path, flux_kontext_path):
+        self.pose2d = Pose2d(checkpoint=pose2d_checkpoint_path, detector_checkpoint=det_checkpoint_path)
+
+        model_cfg = "sam2_hiera_l.yaml"
+        if sam_checkpoint_path is not None:
+            self.predictor = build_sam2_video_predictor(model_cfg, sam_checkpoint_path)    
+        if flux_kontext_path is not None:
+            self.flux_kontext = FluxKontextPipeline.from_pretrained(flux_kontext_path, torch_dtype=torch.bfloat16).to("cuda")
+
+    def __call__(self, video_path, refer_image_path, output_path, resolution_area=[1280, 720], fps=30, iterations=3, k=7, w_len=1, h_len=1, retarget_flag=False, use_flux=False, replace_flag=False):
+        if replace_flag:
+
+            video_reader = VideoReader(video_path)
+            frame_num = len(video_reader)
+            print('frame_num: {}'.format(frame_num))
+            
+            video_fps = video_reader.get_avg_fps()
+            print('video_fps: {}'.format(video_fps))
+            print('fps: {}'.format(fps))
+
+            # TODO: Maybe we can switch to PyAV later, which can get accurate frame num
+            duration = video_reader.get_frame_timestamp(-1)[-1]      
+            expected_frame_num = int(duration * video_fps + 0.5) 
+            ratio = abs((frame_num - expected_frame_num)/frame_num)         
+            if ratio > 0.1:
+                print("Warning: The difference between the actual number of frames and the expected number of frames is two large")
+                frame_num = expected_frame_num
+
+            if fps == -1:
+                fps = video_fps
+
+            target_num = int(frame_num / video_fps * fps)
+            print('target_num: {}'.format(target_num))
+            idxs = get_frame_indices(frame_num, video_fps, target_num, fps)
+            frames = video_reader.get_batch(idxs).asnumpy()
+
+            frames = [resize_by_area(frame, resolution_area[0] * resolution_area[1], divisor=16) for frame in frames]
+            height, width = frames[0].shape[:2]
+            logger.info(f"Processing pose meta")
+
+
+            tpl_pose_metas = self.pose2d(frames)
+
+            face_images = []
+            for idx, meta in enumerate(tpl_pose_metas):
+                face_bbox_for_image = get_face_bboxes(meta['keypoints_face'][:, :2], scale=1.3,
+                                                    image_shape=(frames[0].shape[0], frames[0].shape[1]))
+
+                x1, x2, y1, y2 = face_bbox_for_image
+                face_image = frames[idx][y1:y2, x1:x2]
+                face_image = cv2.resize(face_image, (512, 512))
+                face_images.append(face_image)
+
+            logger.info(f"Processing reference image: {refer_image_path}")
+            refer_img = cv2.imread(refer_image_path)
+            src_ref_path = os.path.join(output_path, 'src_ref.png')
+            shutil.copy(refer_image_path, src_ref_path)
+            refer_img = refer_img[..., ::-1]
+
+            refer_img = padding_resize(refer_img, height, width)
+            logger.info(f"Processing template video: {video_path}")
+            tpl_retarget_pose_metas = [AAPoseMeta.from_humanapi_meta(meta) for meta in tpl_pose_metas]
+            cond_images = []
+
+            for idx, meta in enumerate(tpl_retarget_pose_metas):
+                canvas = np.zeros_like(refer_img)
+                conditioning_image = draw_aapose_by_meta_new(canvas, meta)
+                cond_images.append(conditioning_image)
+            masks = self.get_mask_from_face_bbox(frames, 400, tpl_pose_metas)
+
+            bg_images = []
+            aug_masks = []
+
+            for frame, mask in zip(frames, masks):
+                if iterations > 0:
+                    _, each_mask = get_mask_body_img(frame, mask, iterations=iterations, k=k)
+                    each_aug_mask = get_aug_mask(each_mask, w_len=w_len, h_len=h_len)
+                else:
+                    each_aug_mask = mask
+
+                each_bg_image = frame * (1 - each_aug_mask[:, :, None])
+                bg_images.append(each_bg_image)
+                aug_masks.append(each_aug_mask)
+
+            src_face_path = os.path.join(output_path, 'src_face.mp4')
+            mpy.ImageSequenceClip(face_images, fps=fps).write_videofile(src_face_path)
+
+            src_pose_path = os.path.join(output_path, 'src_pose.mp4')
+            mpy.ImageSequenceClip(cond_images, fps=fps).write_videofile(src_pose_path)
+
+            src_bg_path = os.path.join(output_path, 'src_bg.mp4')
+            mpy.ImageSequenceClip(bg_images, fps=fps).write_videofile(src_bg_path)
+
+            aug_masks_new = [np.stack([mask * 255, mask * 255, mask * 255], axis=2) for mask in aug_masks]
+            src_mask_path = os.path.join(output_path, 'src_mask.mp4')
+            mpy.ImageSequenceClip(aug_masks_new, fps=fps).write_videofile(src_mask_path)
+            return True
+        else:
+            logger.info(f"Processing reference image: {refer_image_path}")
+            refer_img = cv2.imread(refer_image_path)
+            src_ref_path = os.path.join(output_path, 'src_ref.png')
+            shutil.copy(refer_image_path, src_ref_path)
+            refer_img = refer_img[..., ::-1]
+            
+            refer_img = resize_by_area(refer_img, resolution_area[0] * resolution_area[1], divisor=16)
+            
+            refer_pose_meta = self.pose2d([refer_img])[0]
+
+
+            logger.info(f"Processing template video: {video_path}")
+            video_reader = VideoReader(video_path)
+            frame_num = len(video_reader)
+            print('frame_num: {}'.format(frame_num))
+
+            video_fps = video_reader.get_avg_fps()
+            print('video_fps: {}'.format(video_fps))
+            print('fps: {}'.format(fps))
+
+            # TODO: Maybe we can switch to PyAV later, which can get accurate frame num
+            duration = video_reader.get_frame_timestamp(-1)[-1]      
+            expected_frame_num = int(duration * video_fps + 0.5) 
+            ratio = abs((frame_num - expected_frame_num)/frame_num)         
+            if ratio > 0.1:
+                print("Warning: The difference between the actual number of frames and the expected number of frames is two large")
+                frame_num = expected_frame_num
+
+            if fps == -1:
+                fps = video_fps
+                
+            target_num = int(frame_num / video_fps * fps)
+            print('target_num: {}'.format(target_num))
+            idxs = get_frame_indices(frame_num, video_fps, target_num, fps)
+            frames = video_reader.get_batch(idxs).asnumpy()
+
+            logger.info(f"Processing pose meta")
+
+            tpl_pose_meta0 = self.pose2d(frames[:1])[0]
+            tpl_pose_metas = self.pose2d(frames)
+
+            face_images = []
+            for idx, meta in enumerate(tpl_pose_metas):
+                face_bbox_for_image = get_face_bboxes(meta['keypoints_face'][:, :2], scale=1.3,
+                                                    image_shape=(frames[0].shape[0], frames[0].shape[1]))
+
+                x1, x2, y1, y2 = face_bbox_for_image
+                face_image = frames[idx][y1:y2, x1:x2]
+                face_image = cv2.resize(face_image, (512, 512))
+                face_images.append(face_image)
+
+            if retarget_flag:
+                if use_flux:
+                    tpl_prompt, refer_prompt = self.get_editing_prompts(tpl_pose_metas, refer_pose_meta)
+                    refer_input = Image.fromarray(refer_img)
+                    refer_edit = self.flux_kontext(
+                            image=refer_input,
+                            height=refer_img.shape[0],
+                            width=refer_img.shape[1],
+                            prompt=refer_prompt,
+                            guidance_scale=2.5,
+                            num_inference_steps=28,
+                        ).images[0]
+                    
+                    refer_edit = Image.fromarray(padding_resize(np.array(refer_edit), refer_img.shape[0], refer_img.shape[1]))
+                    refer_edit_path = os.path.join(output_path, 'refer_edit.png')
+                    refer_edit.save(refer_edit_path)
+                    refer_edit_pose_meta = self.pose2d([np.array(refer_edit)])[0]
+
+                    tpl_img = frames[1]
+                    tpl_input = Image.fromarray(tpl_img)
+                    
+                    tpl_edit = self.flux_kontext(
+                            image=tpl_input,
+                            height=tpl_img.shape[0],
+                            width=tpl_img.shape[1],
+                            prompt=tpl_prompt,
+                            guidance_scale=2.5,
+                            num_inference_steps=28,
+                        ).images[0]
+                    
+                    tpl_edit = Image.fromarray(padding_resize(np.array(tpl_edit), tpl_img.shape[0], tpl_img.shape[1]))
+                    tpl_edit_path = os.path.join(output_path, 'tpl_edit.png')
+                    tpl_edit.save(tpl_edit_path)
+                    tpl_edit_pose_meta0 = self.pose2d([np.array(tpl_edit)])[0]
+                    tpl_retarget_pose_metas = get_retarget_pose(tpl_pose_meta0, refer_pose_meta, tpl_pose_metas, tpl_edit_pose_meta0, refer_edit_pose_meta)
+                else:
+                    tpl_retarget_pose_metas = get_retarget_pose(tpl_pose_meta0, refer_pose_meta, tpl_pose_metas, None, None)
+            else:
+               tpl_retarget_pose_metas = [AAPoseMeta.from_humanapi_meta(meta) for meta in tpl_pose_metas]
+
+            cond_images = []
+            for idx, meta in enumerate(tpl_retarget_pose_metas):
+                if retarget_flag:
+                    canvas = np.zeros_like(refer_img)
+                    conditioning_image = draw_aapose_by_meta_new(canvas, meta)
+                else:
+                    canvas = np.zeros_like(frames[0])
+                    conditioning_image = draw_aapose_by_meta_new(canvas, meta)
+                    conditioning_image = padding_resize(conditioning_image, refer_img.shape[0], refer_img.shape[1])
+
+                cond_images.append(conditioning_image)
+
+            src_face_path = os.path.join(output_path, 'src_face.mp4')
+            mpy.ImageSequenceClip(face_images, fps=fps).write_videofile(src_face_path)
+
+            src_pose_path = os.path.join(output_path, 'src_pose.mp4')
+            mpy.ImageSequenceClip(cond_images, fps=fps).write_videofile(src_pose_path)
+            return True
+
+    def get_editing_prompts(self, tpl_pose_metas, refer_pose_meta):
+        arm_visible = False
+        leg_visible = False
+        for tpl_pose_meta in tpl_pose_metas:
+            tpl_keypoints = tpl_pose_meta['keypoints_body']
+            if tpl_keypoints[3].all() != 0 or tpl_keypoints[4].all() != 0 or tpl_keypoints[6].all() != 0 or tpl_keypoints[7].all() != 0:
+                if (tpl_keypoints[3][0] <= 1 and tpl_keypoints[3][1] <= 1 and tpl_keypoints[3][2] >= 0.75) or (tpl_keypoints[4][0] <= 1 and tpl_keypoints[4][1] <= 1 and tpl_keypoints[4][2] >= 0.75) or \
+                    (tpl_keypoints[6][0] <= 1 and tpl_keypoints[6][1] <= 1 and tpl_keypoints[6][2] >= 0.75) or (tpl_keypoints[7][0] <= 1 and tpl_keypoints[7][1] <= 1 and tpl_keypoints[7][2] >= 0.75):
+                    arm_visible = True
+            if tpl_keypoints[9].all() != 0 or tpl_keypoints[12].all() != 0 or tpl_keypoints[10].all() != 0 or tpl_keypoints[13].all() != 0:
+                if (tpl_keypoints[9][0] <= 1 and tpl_keypoints[9][1] <= 1 and tpl_keypoints[9][2] >= 0.75) or (tpl_keypoints[12][0] <= 1 and tpl_keypoints[12][1] <= 1 and tpl_keypoints[12][2] >= 0.75) or \
+                    (tpl_keypoints[10][0] <= 1 and tpl_keypoints[10][1] <= 1 and tpl_keypoints[10][2] >= 0.75) or (tpl_keypoints[13][0] <= 1 and tpl_keypoints[13][1] <= 1 and tpl_keypoints[13][2] >= 0.75):
+                    leg_visible = True
+            if arm_visible and leg_visible:
+                break
+        
+        if leg_visible:
+            if tpl_pose_meta['width'] > tpl_pose_meta['height']:
+                tpl_prompt = "Change the person to a standard T-pose (facing forward with arms extended). The person is standing. Feet and Hands are visible in the image."
+            else:
+                tpl_prompt = "Change the person to a standard pose with the face oriented forward and arms extending straight down by the sides. The person is standing. Feet and Hands are visible in the image."
+
+            if refer_pose_meta['width'] > refer_pose_meta['height']:
+                refer_prompt = "Change the person to a standard T-pose (facing forward with arms extended). The person is standing. Feet and Hands are visible in the image."
+            else:
+                refer_prompt = "Change the person to a standard pose with the face oriented forward and arms extending straight down by the sides. The person is standing. Feet and Hands are visible in the image."
+        elif arm_visible:
+            if tpl_pose_meta['width'] > tpl_pose_meta['height']:
+                tpl_prompt = "Change the person to a standard T-pose (facing forward with arms extended). Hands are visible in the image."
+            else:
+                tpl_prompt = "Change the person to a standard pose with the face oriented forward and arms extending straight down by the sides. Hands are visible in the image."
+
+            if refer_pose_meta['width'] > refer_pose_meta['height']:
+                refer_prompt = "Change the person to a standard T-pose (facing forward with arms extended). Hands are visible in the image."
+            else:
+                refer_prompt = "Change the person to a standard pose with the face oriented forward and arms extending straight down by the sides. Hands are visible in the image."
+        else:
+            tpl_prompt = "Change the person to face forward."
+            refer_prompt = "Change the person to face forward."
+
+        return tpl_prompt, refer_prompt
+    
+
+    def get_mask(self, frames, th_step, kp2ds_all):
+        frame_num = len(frames)
+        if frame_num < th_step:
+            num_step = 1
+        else:
+            num_step = (frame_num + th_step) // th_step
+
+        all_mask = []
+        for index in range(num_step):
+            each_frames = frames[index * th_step:(index + 1) * th_step]
+    
+            kp2ds = kp2ds_all[index * th_step:(index + 1) * th_step]
+            if len(each_frames) > 4:
+                key_frame_num = 4
+            elif 4 >= len(each_frames) > 0:
+                key_frame_num = 1
+            else:
+                continue
+
+            key_frame_step = len(kp2ds) // key_frame_num
+            key_frame_index_list = list(range(0, len(kp2ds), key_frame_step))
+
+            key_points_index = [0, 1, 2, 5, 8, 11, 10, 13]
+            key_frame_body_points_list = []
+            for key_frame_index in key_frame_index_list:
+                keypoints_body_list = []
+                body_key_points = kp2ds[key_frame_index]['keypoints_face']
+                for each_index in key_points_index:
+                    each_keypoint = body_key_points[each_index]
+                    if None is each_keypoint:
+                        continue
+                    keypoints_body_list.append(each_keypoint)
+
+                keypoints_body = np.array(keypoints_body_list)[:, :2]
+                wh = np.array([[kp2ds[0]['width'], kp2ds[0]['height']]])
+                points = (keypoints_body * wh).astype(np.int32)
+                key_frame_body_points_list.append(points)
+
+            inference_state = self.predictor.init_state_v2(frames=each_frames)
+            self.predictor.reset_state(inference_state)
+            ann_obj_id = 1
+            for ann_frame_idx, points in zip(key_frame_index_list, key_frame_body_points_list):
+                labels = np.array([1] * points.shape[0], np.int32)
+                _, out_obj_ids, out_mask_logits = self.predictor.add_new_points(
+                    inference_state=inference_state,
+                    frame_idx=ann_frame_idx,
+                    obj_id=ann_obj_id,
+                    points=points,
+                    labels=labels,
+                )
+
+            video_segments = {}
+            for out_frame_idx, out_obj_ids, out_mask_logits in self.predictor.propagate_in_video(inference_state):
+                video_segments[out_frame_idx] = {
+                    out_obj_id: (out_mask_logits[i] > 0.0).cpu().numpy()
+                    for i, out_obj_id in enumerate(out_obj_ids)
+                }
+
+            for out_frame_idx in range(len(video_segments)):
+                for out_obj_id, out_mask in video_segments[out_frame_idx].items():
+                    out_mask = out_mask[0].astype(np.uint8)
+                    all_mask.append(out_mask)
+
+        return all_mask
+    
+    def convert_list_to_array(self, metas):
+        metas_list = []
+        for meta in metas:
+            for key, value in meta.items():
+                if type(value) is list:
+                    value = np.array(value)
+                meta[key] = value
+            metas_list.append(meta)
+        return metas_list
+
+    def get_mask_from_face_bbox(self, frames, th_step, kp2ds_all):
+        """
+        Build masks using a face bounding box per key frame (derived from keypoints_face),
+        then propagate with SAM2 across each chunk of frames.
+        """
+        H, W = frames[0].shape[:2]
+
+        def _clip_box(x1, y1, x2, y2, W, H):
+            x1 = max(0, min(int(x1), W - 1))
+            x2 = max(0, min(int(x2), W - 1))
+            y1 = max(0, min(int(y1), H - 1))
+            y2 = max(0, min(int(y2), H - 1))
+            if x2 <= x1: x2 = min(W - 1, x1 + 1)
+            if y2 <= y1: y2 = min(H - 1, y1 + 1)
+            return x1, y1, x2, y2
+
+        frame_num = len(frames)
+        if frame_num < th_step:
+            num_step = 1
+        else:
+            num_step = (frame_num + th_step) // th_step
+
+        all_mask = []
+
+        for step_idx in range(num_step):
+            each_frames = frames[step_idx * th_step:(step_idx + 1) * th_step]
+            kp2ds = kp2ds_all[step_idx * th_step:(step_idx + 1) * th_step]
+            if len(each_frames) == 0:
+                continue
+
+            # pick a few key frames in this chunk
+            key_frame_num = 4 if len(each_frames) > 4 else 1
+            key_frame_step = max(1, len(kp2ds) // key_frame_num)
+            key_frame_index_list = list(range(0, len(kp2ds), key_frame_step))[:key_frame_num]
+
+            # compute face boxes on the selected key frames
+            key_frame_boxes = []
+            for kfi in key_frame_index_list:
+                meta = kp2ds[kfi]
+                # get_face_bboxes returns (x1, x2, y1, y2) in your code
+                x1, x2, y1, y2 = get_face_bboxes(
+                    meta['keypoints_face'][:, :2],
+                    scale=1.3,
+                    image_shape=(H, W)
+                )
+                x1, y1, x2, y2 = _clip_box(x1, y1, x2, y2, W, H)
+                key_frame_boxes.append(np.array([x1, y1, x2, y2], dtype=np.float32))
+
+            # init SAM2 for this chunk
+            with torch.no_grad(), torch.cuda.amp.autocast(dtype=torch.float16):
+                inference_state = self.predictor.init_state_v2(frames=each_frames)
+                self.predictor.reset_state(inference_state)
+                ann_obj_id = 1
+
+                # seed with box prompts (preferred), else fall back to points
+                for ann_frame_idx, box_xyxy in zip(key_frame_index_list, key_frame_boxes):
+                    used_box = False
+                    try:
+                        # If your predictor exposes a box API, this is ideal.
+                        _ = self.predictor.add_new_box(
+                            inference_state=inference_state,
+                            frame_idx=ann_frame_idx,
+                            obj_id=ann_obj_id,
+                            box=box_xyxy[None, :]  # shape (1, 4)
+                        )
+                        used_box = True
+                    except Exception:
+                        used_box = False
+
+                    if not used_box:
+                        # Fallback: sample a few positive points inside the box
+                        x1, y1, x2, y2 = box_xyxy.astype(int)
+                        cx, cy = (x1 + x2) // 2, (y1 + y2) // 2
+                        pts = np.array([
+                            [cx, cy],
+                            [x1 + (x2 - x1) // 4, cy],
+                            [x2 - (x2 - x1) // 4, cy],
+                            [cx, y1 + (y2 - y1) // 4],
+                            [cx, y2 - (y2 - y1) // 4],
+                        ], dtype=np.int32)
+                        labels = np.ones(len(pts), dtype=np.int32)  # 1 = positive
+                        _ = self.predictor.add_new_points(
+                            inference_state=inference_state,
+                            frame_idx=ann_frame_idx,
+                            obj_id=ann_obj_id,
+                            points=pts,
+                            labels=labels,
+                        )
+
+                # propagate across the chunk
+                video_segments = {}
+                for out_frame_idx, out_obj_ids, out_mask_logits in self.predictor.propagate_in_video(inference_state):
+                    video_segments[out_frame_idx] = {
+                        out_obj_id: (out_mask_logits[i] > 0.0).cpu().numpy()
+                        for i, out_obj_id in enumerate(out_obj_ids)
+                    }
+
+            # collect masks (single object id)
+            for out_frame_idx in range(len(video_segments)):
+                # (H, W) boolean/uint8
+                mask = next(iter(video_segments[out_frame_idx].values()))
+                mask = mask[0].astype(np.uint8)
+                all_mask.append(mask)
+
+        return all_mask
+    def get_mask_from_face_point(self, frames, th_step, kp2ds_all):
+        """
+        Build masks using a single face *center point* per key frame,
+        then propagate with SAM2 across each chunk of frames.
+        """
+        H, W = frames[0].shape[:2]
+
+        frame_num = len(frames)
+        num_step = 1 if frame_num < th_step else (frame_num + th_step) // th_step
+
+        all_mask = []
+
+        for step_idx in range(num_step):
+            each_frames = frames[step_idx * th_step:(step_idx + 1) * th_step]
+            kp2ds = kp2ds_all[step_idx * th_step:(step_idx + 1) * th_step]
+            if len(each_frames) == 0:
+                continue
+
+            # choose a few key frames to seed the object
+            key_frame_num = 1
+            key_frame_step = max(1, len(kp2ds) // key_frame_num)
+            key_frame_index_list = list(range(0, len(kp2ds), key_frame_step))[:key_frame_num]
+
+            # compute center point from face bbox for each selected key frame
+            center_pts = []
+            for kfi in key_frame_index_list:
+                meta = kp2ds[kfi]
+                # your helper returns (x1, x2, y1, y2)
+                x1, x2, y1, y2 = get_face_bboxes(
+                    meta['keypoints_face'][:, :2],
+                    scale=1.3,
+                    image_shape=(H, W)
+                )
+                cx = (x1 + x2) // 2
+                cy = (y1 + y2) // 2
+                # clip just in case
+                cx = int(max(0, min(cx, W - 1)))
+                cy = int(max(0, min(cy, H - 1)))
+                center_pts.append(np.array([cx, cy], dtype=np.int32))
+
+            # init SAM2 for this chunk
+            inference_state = self.predictor.init_state_v2(frames=each_frames)
+            self.predictor.reset_state(inference_state)
+            ann_obj_id = 1
+
+            # seed each key frame with a single positive point at the face center
+            for ann_frame_idx, pt in zip(key_frame_index_list, center_pts):
+                pts = pt[None, :]  # shape (1, 2)
+                labels = np.ones(1, dtype=np.int32)  # 1 = positive
+                _ = self.predictor.add_new_points(
+                    inference_state=inference_state,
+                    frame_idx=ann_frame_idx,
+                    obj_id=ann_obj_id,
+                    points=pts,
+                    labels=labels,
+                )
+
+            # propagate across the chunk
+            video_segments = {}
+            for out_frame_idx, out_obj_ids, out_mask_logits in self.predictor.propagate_in_video(inference_state):
+                video_segments[out_frame_idx] = {
+                    out_obj_id: (out_mask_logits[i] > 0.0).cpu().numpy()
+                    for i, out_obj_id in enumerate(out_obj_ids)
+                }
+
+            # collect masks (single object id)
+            for out_frame_idx in range(len(video_segments)):
+                mask = next(iter(video_segments[out_frame_idx].values()))
+                mask = mask[0].astype(np.uint8)
+                all_mask.append(mask)
+
+        return all_mask
+
+    def get_face_bbox_masks(self, frames, kp2ds_all, scale=1.3, feather_px=0, keep_soft=False):
+        """
+        Create a per-frame mask that's simply the face bounding box.
+        - scale: bbox scale factor used by get_face_bboxes
+        - feather_px: optional Gaussian blur in pixels to feather edges
+        - keep_soft: if True, keep float [0,1] soft mask (after blur); else binarize to {0,1}
+        """
+        H, W = frames[0].shape[:2]
+
+        def _clip_box(x1, y1, x2, y2):
+            x1 = max(0, min(int(x1), W - 1))
+            x2 = max(0, min(int(x2), W - 1))
+            y1 = max(0, min(int(y1), H - 1))
+            y2 = max(0, min(int(y2), H - 1))
+            if x2 <= x1: x2 = min(W - 1, x1 + 1)
+            if y2 <= y1: y2 = min(H - 1, y1 + 1)
+            return x1, y1, x2, y2
+
+        masks = []
+        last_box = None
+        for meta in kp2ds_all:
+            # get_face_bboxes returns (x1, x2, y1, y2)
+            try:
+                x1, x2, y1, y2 = get_face_bboxes(
+                    meta['keypoints_face'][:, :2],
+                    scale=scale,
+                    image_shape=(H, W)
+                )
+                x1, y1, x2, y2 = _clip_box(x1, y1, x2, y2)
+                last_box = (x1, y1, x2, y2)
+            except Exception:
+                # fallback: reuse last seen box to avoid holes
+                if last_box is None:
+                    # no detection yet: push empty mask
+                    masks.append(np.zeros((H, W), dtype=np.uint8))
+                    continue
+                x1, y1, x2, y2 = last_box
+
+            m = np.zeros((H, W), dtype=np.float32)
+            m[y1:y2, x1:x2] = 1.0
+
+            if feather_px and feather_px > 0:
+                # kernel size must be odd and >= 3
+                k = max(3, int(feather_px) | 1)
+                m = cv2.GaussianBlur(m, (k, k), 0)
+
+            if keep_soft:
+                masks.append(m)  # float [0,1]
+            else:
+                masks.append((m >= 0.5).astype(np.uint8))  # hard {0,1}
+
+        return masks
+