models/wan/scail/__init__.py

from __future__ import annotations

__all__ = ["ScailPoseProcessor"]

import numpy as np
import torch
from PIL import Image, ImageOps

from shared.utils import files_locator as fl
from .nlf import load_multiperson_nlf_eager

from .scail_pose_align3d import solve_new_camera_params_central, solve_new_camera_params_down
from .scail_pose_dwpose import DWposeDetector
from .scail_pose_nlf import (
    collect_smpl_poses,
    intrinsic_matrix_from_field_of_view,
    process_data_to_COCO_format,
    process_video_multi_nlf,
    process_video_nlf,
    recollect_dwposes,
    recollect_nlf,
    render_multi_nlf_as_images,
    render_nlf_as_images,
    scale_faces,
    shift_dwpose_according_to_nlf,
)
from .scail_pose_multi import change_poses_to_limit_num, get_largest_bbox_indices


def _pil_resample(name: str):
    try:
        return getattr(Image.Resampling, name)
    except AttributeError:  # Pillow<9
        return getattr(Image, name)


_RESAMPLE_LANCZOS = _pil_resample("LANCZOS")
_RESAMPLE_NEAREST = _pil_resample("NEAREST")


def _tensor_frame_to_uint8_rgb(frame_chw: torch.Tensor) -> np.ndarray:
    frame = frame_chw.detach()
    if frame.dim() != 3:
        raise ValueError(f"Expected CHW tensor, got shape {tuple(frame.shape)}")

    if frame.shape[0] == 1:
        frame = frame.expand(3, -1, -1)
    if frame.shape[0] >= 4:
        frame = frame[:3]

    frame = frame.to(dtype=torch.float32)
    if frame.numel() == 0:
        return np.zeros((0, 0, 3), dtype=np.uint8)

    if float(frame.min()) < 0.0:
        frame = (frame + 1.0) * 127.5
    elif float(frame.max()) <= 1.0:
        frame = frame * 255.0

    frame = frame.clamp(0, 255).to(torch.uint8)
    return frame.permute(1, 2, 0).cpu().numpy()


def _tensor_mask_to_pil(mask: torch.Tensor) -> Image.Image:
    m = mask.detach()
    if m.dim() == 3:
        m = m[0]
    if m.dim() != 2:
        raise ValueError(f"Expected HW (or 1HW) mask tensor, got shape {tuple(mask.shape)}")

    m = m.to(dtype=torch.float32)
    if m.numel() == 0:
        return Image.fromarray(np.zeros((0, 0), dtype=np.uint8), mode="L")

    if float(m.min()) < 0.0:
        m = (m + 1.0) * 0.5
    elif float(m.max()) > 1.0:
        m = m / 255.0

    m = m.clamp(0.0, 1.0)
    m8 = (m * 255.0).to(torch.uint8).cpu().numpy()
    return Image.fromarray(m8, mode="L")


class ScailPoseProcessor:
    """
    WanGP wrapper around the official SCAIL-Pose preprocessing pipeline:
    - DWpose (YOLOX+RTMPose) for 2D keypoints + bbox.
    - NLFPose for 3D pose.
    - Taichi cylinder renderer + 2D overlay.
    """

    def __init__(self, gpu_id: int = 0, *, multi_person: bool = False, max_people: int = 2):
        if not torch.cuda.is_available():
            raise RuntimeError("SCAIL pose preprocessing requires CUDA (GPU-only).")

        self.gpu_id = int(gpu_id)
        self.multi_person = bool(multi_person)
        self.max_people = max(1, int(max_people))

        # Some upstream helpers use bare `.cuda()` / default device semantics.
        torch.cuda.set_device(self.gpu_id)
        self.detector = DWposeDetector(use_batch=False).to(self.gpu_id)

        eager_ckpt = fl.locate_file("pose/nlf_l_multi_0.3.2.eager.safetensors")
        yolox_onnx = fl.locate_file("pose/yolox_l.onnx")
        self.model_nlf = load_multiperson_nlf_eager(
            checkpoint_path=eager_ckpt,
            yolox_onnx_path=yolox_onnx,
            device=f'cuda:{self.gpu_id}'
        )

        # Lazy-loaded for multi-person mode.
        self._sam_segmenter = None
        self._matanyone_model = None
        self._matanyone_version = None

    def unload(self):
        """Release NLF model from memory (RAM and VRAM)."""
        import gc

        if hasattr(self, 'model_nlf') and self.model_nlf is not None:
            # Move model to CPU first to free VRAM
            self.model_nlf.to('cpu')
            # Delete the model
            del self.model_nlf
            self.model_nlf = None

        if hasattr(self, 'detector') and self.detector is not None:
            self.detector.to('cpu')
            del self.detector
            self.detector = None

        if getattr(self, "_sam_segmenter", None) is not None:
            try:
                self._sam_segmenter.model.to("cpu")
            except Exception:
                pass
            self._sam_segmenter = None

        if getattr(self, "_matanyone_model", None) is not None:
            try:
                self._matanyone_model.to("cpu")
            except Exception:
                pass
            self._matanyone_model = None
            self._matanyone_version = None

        # Force garbage collection and empty CUDA cache
        gc.collect()
        torch.cuda.empty_cache()

    def _ensure_sam_loaded(self) -> None:
        if self._sam_segmenter is not None:
            return
        from preprocessing.matanyone.tools.base_segmenter import BaseSegmenter

        # `BaseSegmenter` loads weights from `ckpts/mask/*safetensors` via `files_locator`.
        self._sam_segmenter = BaseSegmenter(SAM_checkpoint=None, model_type="vit_h", device=f"cuda:{self.gpu_id}")

    def _ensure_matanyone_loaded(self) -> None:
        from preprocessing.matanyone.utils.model_assets import get_selected_matanyone_version, load_selected_matanyone_model

        scail_matanyone_config = {"matanyone_version": "v1"}
        selected_version = get_selected_matanyone_version(scail_matanyone_config)
        if self._matanyone_model is not None and self._matanyone_version == selected_version:
            return
        if self._matanyone_model is not None:
            self._matanyone_model.to("cpu")
            self._matanyone_model = None
            torch.cuda.empty_cache()

        self._matanyone_model, self._matanyone_version, _ = load_selected_matanyone_model(scail_matanyone_config)
        self._matanyone_model = self._matanyone_model.eval()
        self._matanyone_model.to(f"cuda:{self.gpu_id}")

    def _extract_and_render_multi(
        self,
        vr_frames_np: list[np.ndarray],
        poses_list,
        det_results_list,
    ) -> torch.Tensor:
        if not vr_frames_np:
            return torch.empty((0,), dtype=torch.float32)

        first_bboxes = det_results_list[0] if det_results_list else None
        if not first_bboxes:
            return torch.empty((0,), dtype=torch.float32)

        indices = get_largest_bbox_indices(first_bboxes, self.max_people)
        if not indices:
            return torch.empty((0,), dtype=torch.float32)

        height, width = vr_frames_np[0].shape[:2]
        pose0 = poses_list[0]

        # 1) Build SAM prompts from DWpose (bbox + a few stable keypoints).
        considered_points = {0, 1, 14, 15}
        box_list_px: list[np.ndarray] = []
        points_list_px: list[np.ndarray] = []
        for idx in indices:
            x1, y1, x2, y2 = first_bboxes[idx]
            box_px = np.array([x1 * width, y1 * height, x2 * width, y2 * height], dtype=np.float32)

            candidate = np.asarray(pose0["bodies"]["candidate"][idx], dtype=np.float32)
            subset = np.asarray(pose0["bodies"]["subset"][idx])
            subset_mod = subset.copy()
            for k in range(len(subset_mod)):
                if k not in considered_points:
                    subset_mod[k] = -1

            pts = candidate[subset_mod != -1]
            pts = np.asarray(pts, dtype=np.float32).reshape(-1, 2)
            pts = pts[(pts[:, 0] >= 0.0) & (pts[:, 1] >= 0.0)]
            if pts.shape[0] == 0:
                pts = np.array([[(x1 + x2) * 0.5, (y1 + y2) * 0.5]], dtype=np.float32)

            pts_px = pts.copy()
            pts_px[:, 0] *= width
            pts_px[:, 1] *= height

            box_list_px.append(box_px)
            points_list_px.append(pts_px)

        # 2) Initial per-person masks on the first frame (SAM).
        self._ensure_sam_loaded()
        self._sam_segmenter.reset_image()
        self._sam_segmenter.set_image(vr_frames_np[0])

        init_masks: list[np.ndarray] = []
        for box_px, pts_px in zip(box_list_px, points_list_px):
            labels = np.ones((pts_px.shape[0],), dtype=np.int32)
            masks, scores, _logits = self._sam_segmenter.predictor.predict(
                point_coords=pts_px,
                point_labels=labels,
                box=box_px,
                multimask_output=True,
            )
            if masks is None or len(masks) == 0:
                init_masks.append(np.zeros((height, width), dtype=np.uint8))
                continue
            best = int(np.argmax(scores)) if scores is not None and len(scores) > 0 else 0
            init_masks.append(masks[best].astype(np.uint8) * 255)

        # Unload SAM to free VRAM before MatAnyone
        if self._sam_segmenter is not None:
            self._sam_segmenter.reset_image()
            self._sam_segmenter.model.to("cpu")
            del self._sam_segmenter
            self._sam_segmenter = None
            torch.cuda.empty_cache()

        # 3) Track/segment each person across the video (MatAnyone).
        self._ensure_matanyone_loaded()
        from preprocessing.matanyone.matanyone.inference.inference_core import InferenceCore
        from preprocessing.matanyone.matanyone_wrapper import matanyone as matanyone_run

        vr_frames_list = []
        for mask0 in init_masks:
            processor = InferenceCore(self._matanyone_model, cfg=self._matanyone_model.cfg)
            masked_frames, _alpha = matanyone_run(processor, vr_frames_np, mask0)
            vr_frames_list.append(torch.from_numpy(np.stack(masked_frames, axis=0)))
            del processor

        # Unload MatAnyone to free VRAM before NLF
        if self._matanyone_model is not None:
            self._matanyone_model.to("cpu")
            del self._matanyone_model
            self._matanyone_model = None
            self._matanyone_version = None
            torch.cuda.empty_cache()

        # 4) Run NLF per segmented person-video.
        nlf_results = process_video_multi_nlf(self.model_nlf, vr_frames_list)

        # 5) Limit DWpose overlays to the same number of persons for visualization.
        poses_vis, det_vis = change_poses_to_limit_num(list(poses_list), list(det_results_list), num_bboxes=len(indices))
        frames_np_rgba = render_multi_nlf_as_images(nlf_results, poses_vis, reshape_pool=None, intrinsic_matrix=None)

        frames_rgb = np.stack([f[:, :, :3] for f in frames_np_rgba], axis=0).astype(np.float32)
        frames_t = torch.from_numpy(frames_rgb).permute(0, 3, 1, 2)  # T,C,H,W
        frames_t = frames_t / 127.5 - 1.0
        return frames_t.permute(1, 0, 2, 3).contiguous()

    def extract_and_render(
        self,
        video_frames: torch.Tensor,
        ref_image: Image.Image,
        mask_frames: torch.Tensor | None = None,
        align_pose: bool = True,
    ) -> torch.Tensor:
        torch.cuda.set_device(self.gpu_id)

        if video_frames is None or not isinstance(video_frames, torch.Tensor):
            raise TypeError("video_frames must be a torch.Tensor")
        if video_frames.dim() != 4:
            raise ValueError(f"Expected video tensor (C,T,H,W), got {tuple(video_frames.shape)}")

        if ref_image is None or not isinstance(ref_image, Image.Image):
            raise TypeError("ref_image must be a PIL.Image")

        out_w, out_h = ref_image.size
        target_size = (int(out_w), int(out_h))

        c, t, h, w = video_frames.shape
        if t == 0:
            return torch.empty((0,), dtype=torch.float32)

        pil_frames: list[Image.Image] = []
        vr_frames_np: list[np.ndarray] = []

        for frame_idx in range(t):
            frame_np = _tensor_frame_to_uint8_rgb(video_frames[:, frame_idx])
            frame_pil = Image.fromarray(frame_np, mode="RGB")

            # Output dimensions must match the ref image; also ensures internal pose extraction
            # resolution is <= output resolution.
            frame_pil = ImageOps.fit(frame_pil, target_size, method=_RESAMPLE_LANCZOS, centering=(0.5, 0.5))

            if mask_frames is not None:
                mask_pil = _tensor_mask_to_pil(mask_frames[:, frame_idx] if mask_frames.dim() == 4 else mask_frames[frame_idx])
                mask_pil = ImageOps.fit(mask_pil, target_size, method=_RESAMPLE_NEAREST, centering=(0.5, 0.5))
                mask_arr = np.array(mask_pil, dtype=np.float32) / 255.0
                rgb_arr = np.array(frame_pil, dtype=np.float32)
                rgb_arr = (rgb_arr * mask_arr[:, :, None]).clip(0, 255).astype(np.uint8)
                frame_pil = Image.fromarray(rgb_arr, mode="RGB")

            pil_frames.append(frame_pil)
            vr_frames_np.append(np.array(frame_pil, dtype=np.uint8))

        detector_return_list = [self.detector(pil_frame) for pil_frame in pil_frames]
        poses, _scores, det_results = zip(*detector_return_list)

        if self.multi_person:
            return self._extract_and_render_multi(vr_frames_np, list(poses), list(det_results))

        vr_frames = torch.from_numpy(np.stack(vr_frames_np, axis=0))  # (T,H,W,3) uint8
        nlf_results = process_video_nlf(self.model_nlf, vr_frames, det_results)

        # If nothing was detected across the whole clip, return empty tensor (caller handles it).
        first_pose_idx = None
        for i, item in enumerate(nlf_results):
            if len(item.get("nlfpose", [])) == 0:
                continue
            first = item["nlfpose"][0]
            if first is None:
                continue
            try:
                has_any = len(first) > 0
            except TypeError:
                has_any = False
            if has_any:
                first_pose_idx = i
                break
        if first_pose_idx is None:
            return torch.empty((0,), dtype=torch.float32)

        target_H, target_W = out_h, out_w
        ori_camera_pose = intrinsic_matrix_from_field_of_view([target_H, target_W])
        ori_focal = ori_camera_pose[0, 0]

        if align_pose:
            # Reference image pose + NLF (single frame)
            ref_rgb = ref_image.convert("RGB")
            ref_rgb = ImageOps.fit(ref_rgb, target_size, method=_RESAMPLE_LANCZOS, centering=(0.5, 0.5))
            pose_ref, _score_ref, det_result_ref = self.detector(ref_rgb)
            if det_result_ref is None or len(det_result_ref) == 0:
                align_pose = False
            else:
                vr_ref = torch.from_numpy(np.array(ref_rgb, dtype=np.uint8)).unsqueeze(0)
                nlf_results_ref = process_video_nlf(self.model_nlf, vr_ref, [det_result_ref])

                pose_3d_first_driving_frame = nlf_results[first_pose_idx]["nlfpose"][0][0].cpu().numpy()
                pose_3d_coco_first_driving_frame = process_data_to_COCO_format(pose_3d_first_driving_frame)

                poses_2d_ref = pose_ref["bodies"]["candidate"][0][:14]
                poses_2d_ref[:, 0] = poses_2d_ref[:, 0] * target_W
                poses_2d_ref[:, 1] = poses_2d_ref[:, 1] * target_H

                poses_2d_subset = pose_ref["bodies"]["subset"][0][:14]
                pose_3d_coco_first_driving_frame = pose_3d_coco_first_driving_frame[:14]

                valid_upper_indices = []
                valid_lower_indices = []
                upper_body_indices = [0, 2, 3, 5, 6]
                lower_body_indices = [9, 10, 12, 13]
                for j in range(len(poses_2d_subset)):
                    if poses_2d_subset[j] != -1.0 and np.sum(pose_3d_coco_first_driving_frame[j]) != 0:
                        if j in upper_body_indices:
                            valid_upper_indices.append(j)
                        if j in lower_body_indices:
                            valid_lower_indices.append(j)

                if len(valid_lower_indices) >= 4:
                    valid_indices = [1] + valid_lower_indices
                    new_camera_intrinsics, scale_m = solve_new_camera_params_down(
                        pose_3d_coco_first_driving_frame[valid_indices], ori_focal, [target_H, target_W], poses_2d_ref[valid_indices]
                    )
                else:
                    valid_indices = [1] + valid_upper_indices
                    new_camera_intrinsics, scale_m = solve_new_camera_params_central(
                        pose_3d_coco_first_driving_frame[valid_indices], ori_focal, [target_H, target_W], poses_2d_ref[valid_indices]
                    )

                poses_list = list(poses)
                _ = scale_faces(poses_list, [pose_ref])
                nlf_results = recollect_nlf(nlf_results)
                poses_list = recollect_dwposes(poses_list)
                shift_dwpose_according_to_nlf(
                    collect_smpl_poses(nlf_results), poses_list, ori_camera_pose, new_camera_intrinsics, target_H, target_W
                )
                frames_np_rgba = render_nlf_as_images(nlf_results, poses_list, reshape_pool=None, intrinsic_matrix=new_camera_intrinsics)

                frames_rgb = np.stack([f[:, :, :3] for f in frames_np_rgba], axis=0).astype(np.float32)
                frames_t = torch.from_numpy(frames_rgb).permute(0, 3, 1, 2)  # T,C,H,W
                frames_t = frames_t / 127.5 - 1.0
                return frames_t.permute(1, 0, 2, 3).contiguous()

        # Non-aligned path (official fallback)
        nlf_results = recollect_nlf(nlf_results)
        frames_np_rgba = render_nlf_as_images(nlf_results, list(poses), reshape_pool=None, intrinsic_matrix=ori_camera_pose)
        frames_rgb = np.stack([f[:, :, :3] for f in frames_np_rgba], axis=0).astype(np.float32)
        frames_t = torch.from_numpy(frames_rgb).permute(0, 3, 1, 2)
        frames_t = frames_t / 127.5 - 1.0
        return frames_t.permute(1, 0, 2, 3).contiguous()