from PIL import Image
from kaggle_gpu_server.engine.plugin_base import ModelPlugin, PluginCapability, ModelCategory

class DWPosePlugin(ModelPlugin):
    name = "dwpose"
    model_id = "yzd-v/DWPose"
    capability = PluginCapability.POSE_ESTIMATION
    category = ModelCategory.LIGHTWEIGHT
    vram_estimate_mb = 300
    version = "1.0.0"
    description = "Estimates pose using DWPose and returns skeleton visualization"

    def __init__(self):
        super().__init__()
        self.detector = None

    def load(self) -> bool:
        if self._loaded:
            return True
        try:
            from controlnet_aux import DWposeDetector
            self.detector = DWposeDetector.from_pretrained(
                "lllyasviel/ControlNet-v1-1-nightly",
                subfolder="annotator/ckpts"
            )
            # DWPose loads on GPU by default if cuda available
            self._loaded = True
            return True
        except Exception as e:
            print(f"❌ Failed to load DWPose: {e}")
            return False

    def unload(self) -> None:
        if not self._loaded:
            return
        self.detector = None
        self._loaded = False
        import torch
        import gc
        gc.collect()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()

    def _execute(self, inputs: dict) -> dict:
        image = inputs.get("image")
        if image is None:
            raise ValueError("Input 'image' is required")
            
        if self.detector is None:
            # Return empty skeleton
            empty_pose = Image.new("RGBA", image.size, (0, 0, 0, 0))
            return {
                "pose_keypoints": [],
                "pose_image": empty_pose,
                "image": image
            }
            
        # Run pose detection
        # DWPose detector returns a PIL Image of the pose skeleton
        pose_img = self.detector(image)
        
        return {
            "pose_keypoints": [],
            "pose_image": pose_img,
            "image": image
        }