Merge branch 'main' into copilot/replicate-comfyui-sdpose

.gitignore

@@ -31,6 +31,7 @@ Thumbs.db
 *.pth
 *.onnx
 *.safetensors
+!empty_text_encoder/empty_embedding.safetensors
 
 # Output files
 output/

README.md

@@ -111,51 +111,33 @@ Runs pose estimation on one or more images.
 
 ### `IBB Pose — JSON to Image`
 
-Renders an OpenPose JSON string to a pose visualisation image.
+## 模型目录
 
----
+- SDPose 模型：`ComfyUI/models/IBB_POSE/`
+- YOLO 模型：`ComfyUI/models/yolo/`
+- GroundingDINO 模型：`ComfyUI/models/grounding-dino/`
 
-## JSON Format
+首次运行时会自动下载以下资源：
 
-Output follows the standard OpenPose JSON format:
+- `teemosliang/SDPose-Body`
+- `teemosliang/SDPose-Wholebody`
+- GroundingDINO 配置和权重
 
-```json
-{
-  "people": [
-    {
-      "pose_keypoints_2d":       [x, y, score, ...],
-      "face_keypoints_2d":       [...],
-      "hand_left_keypoints_2d":  [...],
-      "hand_right_keypoints_2d": [...],
-      "foot_keypoints_2d":       [...]
-    }
-  ],
-  "canvas_width":  512,
-  "canvas_height": 768
-}
-```
+仓库已内置 `empty_text_encoder/empty_embedding.safetensors`，也保留了 `generate_empty_embedding.py` 以便重新生成。
 
-Compatible with [ComfyUI-OpenPose-Editor](https://github.com/judian17/ComfyUI-OpenPose-Editor-jd).
+## 工作流
 
----
+示例工作流位于：
 
-## Skeleton Keypoint Layout
+- `/workflow/ibb_pose.json`
 
-### Body / OpenPose (18 pts)
-OpenPose-18: nose · neck · shoulders · elbows · wrists · hips · knees · ankles · eyes · ears
+## 兼容说明
 
-### WholeBody (133 pts, DWPose)
-- 0–16: Body (COCO-17)
-- 17–22: Feet (6)
-- 23–90: Face (68)
-- 91–111: Left hand (21)
-- 112–132: Right hand (21)
+- 依赖 ComfyUI 的 `folder_paths` / `model_management`
+- Florence2 检测为可选能力，需要额外安装对应 ComfyUI 节点
+- OpenPose 编辑 JSON 可与 [ComfyUI-OpenPose-Editor-jd](https://github.com/judian17/ComfyUI-OpenPose-Editor-jd) 配合使用
 
----
+## 致谢
 
-## Recommended Workflow
-
-```
-Load Image → IBB Pose Load Model → IBB Pose Run Estimation → Preview Image
-                                                            └→ Save Text (JSON)
-```
+- 上游实现：`judian17/ComfyUI-SDPose-OOD`
+- 原始项目：`T-S-Liang/SDPose-OOD`

__init__.py

@@ -1,7 +1,3 @@
-"""
-IBB_POSE - ComfyUI Pose Estimation Node
-Compatible with Python 3.13, torch 2.10, CUDA 13.x
-"""
 from .ibb_pose_nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
 
 __all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]

eval/SDPose_Evaluator.py

@@ -0,0 +1,140 @@
+# Author: T. S. Liang @ Rama Alpaca
+# Emails: tsliang2001@gmail.com | shuangliang@ramaalpaca.com | sliang57@connect.hku.hk
+# Date: Feb. 2025
+# Description: Evaluation protocol for fine-tuned Stable Diffusion models on pose estimation benchmarks.
+# License: MIT License (see LICENSE for details)
+
+import os
+import torch
+from tqdm import tqdm
+from safetensors.torch import load_file
+
+from accelerate import Accelerator
+from accelerate.logging import get_logger
+from packaging import version
+from tqdm.auto import tqdm
+
+from transformers import CLIPTextModel, CLIPTokenizer
+from diffusers.utils.import_utils import is_xformers_available
+
+from diffusers import (
+    UNet2DConditionModel,
+    AutoencoderKL,
+    DDPMScheduler
+)
+
+from pipelines.SDPose_D_Pipeline import SDPose_D_Pipeline
+
+from models.HeatmapHead import get_heatmap_head
+from models.ModifiedUNet import Modified_forward
+
+from utils.EvalMetric import get_metric
+from utils.dataset_loader import prepare_dataset
+
+
+class SDPose_Evaluator:
+    
+    def __init__(
+        self,
+        args
+        ):
+        
+        self.logger = get_logger(__name__, log_level="INFO")
+        # -------------------- Device --------------------
+        
+        self.args = args
+        self.accelerator = Accelerator()
+
+        unet = UNet2DConditionModel.from_pretrained(
+            self.args.checkpoint_path, subfolder="unet", revision=None,
+            class_embed_type="projection", projection_class_embeddings_input_dim=4,
+            low_cpu_mem_usage=False, device_map=None,
+        )
+
+        unet = Modified_forward(unet, keypoint_scheme = self.args.keypoint_scheme)
+        
+        vae = AutoencoderKL.from_pretrained(self.args.checkpoint_path,
+                                            subfolder='vae')
+
+        tokenizer = CLIPTokenizer.from_pretrained(self.args.checkpoint_path, subfolder='tokenizer')
+        text_encoder = CLIPTextModel.from_pretrained(self.args.checkpoint_path, subfolder='text_encoder')
+
+        # Loading Heatmap Autoencoder
+        dec_path = os.path.join(self.args.checkpoint_path, "decoder", "decoder.safetensors")
+        hm_decoder = get_heatmap_head(mode = self.args.keypoint_scheme)
+        hm_decoder.load_state_dict(load_file(dec_path, device="cpu"), strict=True)
+
+        noise_scheduler = DDPMScheduler.from_pretrained(
+            self.args.checkpoint_path, subfolder='scheduler'
+        )
+        
+        self.model = SDPose_D_Pipeline(unet = unet, vae = vae, tokenizer = tokenizer, text_encoder = text_encoder, decoder = hm_decoder, scheduler = noise_scheduler)
+        
+        # using xformers for efficient attentions.
+        if self.args.enable_xformers_memory_efficient_attention:
+            if is_xformers_available():
+                import xformers
+                xformers_version = version.parse(xformers.__version__)
+                if xformers_version == version.parse("0.0.16"):
+                    self.logger.warn(
+                        "xFormers 0.0.16 cannot be used for training in some GPUs. If you observe problems during training, please update xFormers to at least 0.0.17. See https://huggingface.co/docs/diffusers/main/en/optimization/xformers for more details."
+                    )
+                self.model.unet.enable_xformers_memory_efficient_attention()
+            else:
+                raise ValueError("xformers is not available. Make sure it is installed correctly")
+
+        self.model.vae.requires_grad_(False)
+        self.model.text_encoder.requires_grad_(False)
+        self.model.decoder.requires_grad_(False)
+        self.model.unet.requires_grad_(False)
+
+        self.model.vae.eval()
+        self.model.text_encoder.eval()
+        self.model.decoder.eval()
+        self.model.unet.eval()
+
+    @torch.no_grad()
+    def evaluate(self):
+
+        val_loader = prepare_dataset(batch_size = self.args.eval_batch_size, dataset_name = self.args.dataset_name, dataset_root = self.args.dataset_root, mode = "val", num_workers = self.args.dataloader_num_workers)
+
+
+        self.model.unet, self.model.vae, self.model.decoder, val_loader, self.model.text_encoder = self.accelerator.prepare(self.model.unet, self.model.vae, self.model.decoder, val_loader, self.model.text_encoder)
+        
+        device = self.accelerator.device
+
+        metric = get_metric(ann_file = self.args.ann_file, mode = self.args.keypoint_scheme)
+        
+        metric.dataset_meta = val_loader.dataset.metainfo
+        
+        iterator = tqdm(val_loader, desc="Validating", leave=False) if self.accelerator.is_main_process else val_loader
+        
+        for batch in iterator:
+
+            rgb_in = batch["inputs"] # [B, 3, H, W]
+            rgb_in = rgb_in[:, [2, 1, 0], ...]
+            rgb_in = rgb_in.float() / 255.0
+
+            rgb_in = rgb_in * 2.0 - 1.0
+
+            rgb_in = rgb_in.to(device)  # [B, 3, H, W]
+
+            ds_list = batch['data_samples']
+
+            # Enable flip test by setting test_cfg
+            test_cfg = {'flip_test': True}  # Re-enabled to test our fixes
+            
+            # Run inference with flip test enabled
+            timesteps = [self.args.timestep]
+            outputs = self.model(rgb_in, mode = "predict", data_samples=batch['data_samples'], timesteps = timesteps, test_cfg=test_cfg)
+            
+            samples_for_metric = [s.to_dict() for s in outputs]
+            
+            metric.process(data_samples = samples_for_metric, data_batch = batch)
+
+        self.accelerator.wait_for_everyone()
+        res = metric.evaluate(len(val_loader.dataset))
+
+        if self.accelerator.is_main_process:
+                
+            print(res)

eval/launch_evaluator.py

@@ -0,0 +1,82 @@
+# Author: T. S. Liang @ Rama Alpaca
+# Emails: tsliang2001@gmail.com | shuangliang@ramaalpaca.com | sliang57@connect.hku.hk
+# Date: Feb. 2025
+# License: MIT License (see LICENSE for details)
+
+from SDPose_Evaluator import SDPose_Evaluator
+import argparse
+
+def parse_args():
+
+    parser = argparse.ArgumentParser(description="Evaluate SDPose on Pose Estimation Benchmarks")
+
+    parser.add_argument(
+        "--dataset_name",
+        type=str,
+        default=None,
+        required=True,
+        help="The dataset name for the evaluation benchmark.",
+    )
+
+    parser.add_argument(
+        "--keypoint_scheme",
+        type=str,
+        default="body",
+        required=True,
+        help="The keypoint scheme for the pose estimation model, please fill in body (COCO-17) or wholebody (COCOWholebody-133).",
+    )
+
+    parser.add_argument(
+        "--dataset_root",
+        type=str,
+        default=None,
+        required=True,
+        help="The root path for the datasets.",
+    )
+
+    parser.add_argument(
+        "--checkpoint_path",
+        type=str,
+        default=None,
+        required=True,
+        help="Path to the SDPose checkpoint.",
+    )
+
+    parser.add_argument(
+        "--eval_batch_size", type=int, default=1, help="Batch size (per device) for the evaluation dataloader."
+    )
+
+    parser.add_argument(
+        "--ann_file", type=str, help="The path for the annotation file."
+    )
+
+    parser.add_argument(
+        "--timestep",
+        type=int,
+        default=999,
+    )
+    
+    # dataloaderes
+    parser.add_argument(
+        "--dataloader_num_workers",
+        type=int,
+        default=0,
+        help=(
+            "Number of subprocesses to use for data loading. 0 means that the data will be loaded in the main process."
+        ),
+    )
+    
+    # using xformers for efficient training and inference
+    parser.add_argument(
+        "--enable_xformers_memory_efficient_attention", action="store_true", help="Whether or not to use xformers."
+    )
+
+    args = parser.parse_args()
+
+    return args
+
+
+args = parse_args()
+Trainer = SDPose_Evaluator(args)
+
+Trainer.evaluate()

generate_empty_embedding.py

@@ -0,0 +1,67 @@
+import torch
+from transformers import CLIPTokenizer, CLIPTextModel
+from safetensors.torch import save_file
+import os
+from pathlib import Path
+# No longer need folder_paths for this script
+
+# --- Configuration ---
+# We only need to do this once, so we can hardcode the repo ID.
+# The empty embedding is identical for both Body and WholeBody models.
+REPO_ID = "teemosliang/SDPose-Body"
+# Get the directory of the current script and build the path relative to it
+SCRIPT_DIR = Path(__file__).parent
+OUTPUT_DIR = os.path.join(SCRIPT_DIR, "empty_text_encoder")
+OUTPUT_FILE = os.path.join(OUTPUT_DIR, "empty_embedding.safetensors")
+
+def generate_and_save_empty_embedding():
+    """
+    Loads a CLIP text model, generates an embedding for an empty prompt,
+    and saves it to a safetensors file.
+    """
+    print("--- SDPose Empty Embedding Generator ---")
+
+    # Ensure the output directory exists
+    os.makedirs(OUTPUT_DIR, exist_ok=True)
+    print(f"Output directory prepared: {OUTPUT_DIR}")
+
+    # Check if the file already exists
+    if os.path.exists(OUTPUT_FILE):
+        print(f"'{OUTPUT_FILE}' already exists. Skipping generation.")
+        return
+
+    try:
+        # 1. Load tokenizer and text encoder from Hugging Face
+        print(f"Loading CLIP model from '{REPO_ID}'...")
+        tokenizer = CLIPTokenizer.from_pretrained(REPO_ID, subfolder='tokenizer')
+        text_encoder = CLIPTextModel.from_pretrained(REPO_ID, subfolder='text_encoder')
+        print("Model loaded successfully.")
+
+        # 2. Generate the empty text embedding
+        print("Generating empty text embedding...")
+        prompt = ""
+        text_inputs = tokenizer(
+            prompt,
+            padding="do_not_pad",
+            max_length=tokenizer.model_max_length,
+            truncation=True,
+            return_tensors="pt",
+        )
+        text_input_ids = text_inputs.input_ids
+        empty_text_embed = text_encoder(text_input_ids)[0]
+        print(f"Embedding generated with shape: {empty_text_embed.shape}")
+
+        # 3. Save the embedding to a safetensors file
+        print(f"Saving embedding to '{OUTPUT_FILE}'...")
+        tensor_dict = {"empty_text_embed": empty_text_embed}
+        save_file(tensor_dict, OUTPUT_FILE)
+        print("--- Success! ---")
+        print("Empty embedding has been saved. You no longer need to run this script.")
+
+    except Exception as e:
+        print(f"\n--- An error occurred ---")
+        print(f"Error: {e}")
+        print("Please ensure you have an internet connection and the 'transformers' library is installed.")
+
+if __name__ == "__main__":
+    generate_and_save_empty_embedding()

models/HeatmapHead.py

@@ -0,0 +1,64 @@
+# Author: T. S. Liang @ Rama Alpaca
+# Emails: tsliang2001@gmail.com | shuangliang@ramaalpaca.com | sliang57@connect.hku.hk
+# Date: Jun. 2025
+# Description: SDPose heatmap head factory using mmpose `HeatmapHead` + `UDPHeatmap` codec.
+#   - mode="body"     → 17 keypoints, in_channels=320
+#   - mode="wholebody"→ 133 keypoints, in_channels=640
+#   Defaults: image_size=(768,1024), scale=4 (→ heatmap size = (w/4, h/4)), sigma=6.
+# License: MIT License (see LICENSE for details)
+
+from mmpose.registry import MODELS
+
+def get_heatmap_head(mode = "body"):
+
+    if mode == "body":
+
+        image_size = [768, 1024] ## width x height
+        sigma = 6 ## sigma is 2 for 256
+        scale = 4
+
+        embed_dim = 320
+        num_keypoints = 17
+
+        codec = dict(
+            type='UDPHeatmap', input_size=(image_size[0], image_size[1]), heatmap_size=(int(image_size[0]/scale), int(image_size[1]/scale)), sigma=sigma) ## sigma is 2 for 256
+
+        head_cfg = dict(
+            type='HeatmapHead',
+            in_channels=embed_dim,
+            out_channels=num_keypoints,
+            deconv_out_channels=(320,),
+            deconv_kernel_sizes=(4,),
+            conv_out_channels=(320,),
+            conv_kernel_sizes=(1,),
+            loss=dict(type='KeypointMSELoss', use_target_weight=True),
+            decoder=codec
+        )
+
+    elif mode == "wholebody":
+
+        image_size = [768, 1024] ## width x height
+        sigma = 6 ## sigma is 2 for 256
+        scale = 4
+
+        embed_dim = 640
+        num_keypoints = 133
+
+        codec = dict(
+            type='UDPHeatmap', input_size=(image_size[0], image_size[1]), heatmap_size=(int(image_size[0]/scale), int(image_size[1]/scale)), sigma=sigma) ## sigma is 2 for 256
+
+        head_cfg = dict(
+            type='HeatmapHead',
+            in_channels=embed_dim,
+            out_channels=num_keypoints,
+            deconv_out_channels=(640,),
+            deconv_kernel_sizes=(4,),
+            conv_out_channels=(640,),
+            conv_kernel_sizes=(1,),
+            loss=dict(type='KeypointMSELoss', use_target_weight=True),
+            decoder=codec
+        )
+
+    head = MODELS.build(head_cfg)
+    
+    return head

models/ModifiedUNet.py

@@ -0,0 +1,69 @@
+# Author: T. S. Liang @ Rama Alpaca
+# Emails: tsliang2001@gmail.com | shuangliang@ramaalpaca.com
+# Date: Feb. 2025
+
+"""
+
+Description:
+
+    SDPose UNet Forward Modifier
+
+    Temporarily overrides `UNet.forward` to expose decoder features via forward hooks.
+    - keypoint_scheme: "body" (COCO-17) or "wholebody" (COCO-WholeBody-133)
+    - return_decoder_feats: if True, returns a selected decoder feature (Tensor);
+    otherwise returns the original UNet output.
+"""
+
+def new_forward_kpt17(self, *args, return_decoder_feats=False, **kwargs):
+    self._decoder_feats = []
+
+    def hook_fn(module, input, output):
+        self._decoder_feats.append(output)
+
+    handles = [blk.register_forward_hook(hook_fn) for blk in self.up_blocks]
+
+    out = self._old_forward(*args, **kwargs)
+
+    for h in handles:
+        h.remove()
+
+    if return_decoder_feats:
+
+        feats = self._decoder_feats[::-1]
+        return feats[0]
+    else:
+        return out
+
+def new_forward_kpt133(self, *args, return_decoder_feats=False, **kwargs):
+    self._decoder_feats = []
+
+    def hook_fn(module, input, output):
+        self._decoder_feats.append(output)
+
+    handles = [blk.register_forward_hook(hook_fn) for blk in self.up_blocks]
+
+    out = self._old_forward(*args, **kwargs)
+
+    for h in handles:
+        h.remove()
+
+    if return_decoder_feats:
+
+        feats = self._decoder_feats[::-1]
+        return feats[1]
+    else:
+        return out
+
+def Modified_forward(unet, keypoint_scheme = "body"):
+
+    if keypoint_scheme == "body":
+
+        unet._old_forward = unet.forward
+        unet.forward = new_forward_kpt17.__get__(unet, unet.__class__)
+
+    elif keypoint_scheme == "wholebody":
+
+        unet._old_forward = unet.forward
+        unet.forward = new_forward_kpt133.__get__(unet, unet.__class__)
+
+    return unet

models/__init__.py

@@ -0,0 +1 @@
+# Models package for SDPose OOD

pipelines/SDPose_D_Pipeline.py

@@ -0,0 +1,404 @@
+# Author T. S. Liang @ Rama Alpaca, July. 2024.
+# Email: tsliang2001@gmail.com & shuangliang@ramaalpaca.com
+# Adapted from Marigold: https://github.com/prs-eth/marigold and MMPose: https://github.com/open-mmlab/mmpose
+# License: MIT License (see LICENSE for details)
+
+"""
+    Description:
+    
+    This script is the pipeline utilizing Stable Diffusion Model for the pose perception.
+    Adaptations have been made to ensure full compatibility with MMPose-style
+    decoder heads, enabling seamless integration with existing heatmap-based
+    keypoint decoders and evaluation protocols.
+"""
+
+from typing import Union, Optional
+
+import torch
+import torch.nn as nn
+
+from diffusers import (
+    AutoencoderKL,
+    DiffusionPipeline,
+    UNet2DConditionModel,
+    DDIMScheduler,
+    DDPMScheduler,
+)
+
+from PIL import Image
+
+from tqdm.auto import tqdm
+from transformers import CLIPTextModel, CLIPTokenizer
+
+from itertools import zip_longest
+from mmpose.utils.typing import (ConfigType, InstanceList, OptConfigType,
+                                 OptMultiConfig, PixelDataList, SampleList)
+
+from mmpose.models.heads.heatmap_heads.heatmap_head import HeatmapHead
+    
+class SDPose_D_Pipeline(DiffusionPipeline):
+     
+    """
+    A diffusion-based human pose estimation pipeline that reconstructs pose heatmaps
+    from RGB images using a Stable-Diffusion-style latent denoising process.
+
+    NOTE:
+        This pipeline assumes the scheduler/unet are configured for x0 prediction,
+        i.e., `scheduler.config.prediction_type == "sample"`. In this setting the UNet
+        directly predicts the clean latent x0 at a fixed diffusion step.
+
+    Args:
+        unet (`UNet2DConditionModel`):
+            Conditional UNet used to predict the clean latent `x0`.
+            Inputs: noisy pose/image latents and conditioning (image latent, text embeds).
+        vae (`AutoencoderKL`):
+            Variational Autoencoder for encoding RGB images to latents and for decoding
+            latents (image/heatmap) back to pixel space as needed.
+        text_encoder (`CLIPTextModel`):
+            Text encoder to obtain conditioning embeddings. For text-free pose tasks,
+            empty prompts are typically used.
+        tokenizer (`CLIPTokenizer`):
+            Tokenizer paired with the text encoder (tokenizes empty/neutral prompts).
+        decoder (`HeatmapHead`):
+            A lightweight head mapping denoised latent features to pose heatmaps,
+            typically with shape `[B, num_keypoints, H, W]`.
+        scheduler (`Union[DDIMScheduler, DDPMScheduler]`):
+            Diffusion scheduler defining the diffusion process configuration
+            (only used to set the fixed timestep).  
+            Its `prediction_type` **must be `'sample'`** to ensure x₀ prediction.
+
+    Returns:
+        `SDPose_D_Pipeline`:
+            Initialized pipeline ready for inference/training.
+    """
+
+    def __init__(
+        self,
+        unet: UNet2DConditionModel,
+        vae: AutoencoderKL,
+        text_encoder: CLIPTextModel,
+        tokenizer: CLIPTokenizer,
+        decoder: HeatmapHead,
+        scheduler: Union[DDIMScheduler, DDPMScheduler]
+        ):
+
+        super().__init__()
+        
+        self.register_modules(
+            unet = unet,
+            vae = vae,
+            text_encoder=text_encoder,
+            tokenizer=tokenizer,
+            decoder = decoder,
+            scheduler = scheduler
+        )
+        
+        self.rgb_latent_scale_factor = 0.18215
+
+        self.test_cfg = {
+            'output_keypoint_indices': None,
+            'flip_test': True,
+            'flip_mode': 'heatmap',
+            'shift_heatmap': False
+        }
+
+    @torch.no_grad()
+    def __call__(
+            self,
+            rgb_in: Union[Image.Image, torch.Tensor],
+            show_progress_bar: bool = True,
+            timesteps = None,
+            mode = None,
+            data_samples = None,
+            test_cfg = None
+        ):
+        
+        """
+        Run a forward inference through the SDPose-D pipeline.
+
+        This function performs **x₀ prediction** using a Stable-Diffusion-style UNet
+        conditioned on the RGB image latent, and optionally returns decoded pose heatmaps
+        or full structured pose predictions (MMPose-compatible).
+
+        Args:
+            rgb_in (`PIL.Image.Image` or `torch.Tensor`):
+                Input RGB image or batch tensor of shape `[B, 3, H, W]`.
+                If a PIL image is provided, it should be preprocessed to a normalized tensor
+                compatible with the VAE encoder input range.
+            show_progress_bar (`bool`, optional, defaults to `True`):
+                Whether to display a progress bar during diffusion inference.
+            timesteps (`int` or `list[int]`, optional):
+                Diffusion timestep(s) to use.  
+            mode (`str`, optional):
+                - `"predict"`: Use the decoder's `predict()` method to generate keypoint
+                  predictions and attach them to `data_samples` (for evaluation).  
+                - `None`: Decode the denoised latent directly to a heatmap tensor `[B, K, H, W]`
+                  without MMPose post-processing.
+            data_samples (`SampleList`, optional):
+                A list of MMPose `PoseDataSample` objects required when `mode="predict"`.
+            test_cfg (`dict`, optional):
+                Configuration overrides for test-time behavior.  
+                Examples: `{'flip_test': True, 'shift_heatmap': False}`.
+
+        Returns:
+            - If `mode == "predict"`:
+                `SampleList`:  
+                The input `data_samples` with updated `pred_instances` and optional `pred_fields`
+                containing the decoded keypoints and heatmaps.
+            - Else:
+                `torch.Tensor`:  
+                The predicted pose heatmap tensor of shape `[B, num_keypoints, H, W]`.
+
+        Notes:
+            - This pipeline operates in **x₀ prediction mode**, where the UNet directly predicts
+              the clean latent representation instead of ε or v.  
+            - The pose diffusion process uses a **single, fixed timestep** (default t=999),
+              meaning there is no iterative denoising loop.
+            - `test_cfg` supports `flip_test` inference (horizontal ensemble of original and flipped images).
+            - `scheduler.config.prediction_type` must be set to `"sample"` for correct operation.
+        """
+            
+        # ------------- Predicting pose ---------------
+            
+        # Predict pose images
+
+        self.encode_empty_text()
+        
+        # Update test_cfg if provided
+        if test_cfg is not None:
+            self.test_cfg.update(test_cfg)
+            
+        heatmap_pred = self.single_infer(
+            rgb_in = rgb_in,
+            show_pbar = show_progress_bar,
+            timesteps = timesteps,
+            mode = mode,
+            data_samples = data_samples
+            )
+            
+        return heatmap_pred
+    
+    @torch.no_grad()
+    def single_infer(
+            self,
+            rgb_in: torch.Tensor,
+            timesteps,
+            show_pbar: bool,
+            mode = None,
+            data_samples = None
+        )-> torch.Tensor:
+
+        device = self.unet.device
+        rgb_in = rgb_in.to(device)
+        bsz = rgb_in.shape[0]
+        
+        self.scheduler.set_timesteps(timesteps = timesteps, device=device)
+        
+        timesteps = torch.tensor(timesteps, device=device).long()
+
+        # STEP 1: Encode image
+        if self.test_cfg.get('flip_test', False):
+            # Flip test: process both original and flipped images
+            rgb_in_flipped = torch.flip(rgb_in, dims=[3])  # Flip horizontally (W dimension)
+            
+            # Encode both original and flipped images
+            rgb_latent = self.encode_rgb(rgb_in)
+            rgb_latent_flipped = self.encode_rgb(rgb_in_flipped)
+            
+            # Process both latents
+            feats = []
+            for rgb_latent_curr in [rgb_latent, rgb_latent_flipped]:
+                # Denoising loop for current latent
+                if show_pbar:
+                    iterable = tqdm(
+                        enumerate(timesteps),
+                        total=len(timesteps),
+                        leave=False,
+                        desc=" " * 4 + "Diffusion denoising",
+                    )
+                else:
+                    iterable = enumerate(timesteps)
+
+                self.encode_empty_text()
+                text_embed = self.empty_text_embed.repeat((bsz, 1, 1))
+
+                for i, t in iterable:
+                    task_emb_anno = torch.tensor([1, 0]).float().unsqueeze(0).to(device)
+                    task_emb_anno = torch.cat([torch.sin(task_emb_anno), torch.cos(task_emb_anno)], dim=-1).repeat(bsz, 1)
+
+                    if self.scheduler.config.prediction_type == "sample":
+                        feat = self.unet(
+                            rgb_latent_curr, t, text_embed, class_labels=task_emb_anno, return_dict=False, return_decoder_feats=True
+                        )
+                    else:
+                        raise
+                
+                feats.append((feat,))  # Wrap each feat in a tuple to match expected format
+
+        else:
+            # Original single image processing
+            rgb_latent = self.encode_rgb(rgb_in)
+
+            # Denoising loop
+            if show_pbar:
+                iterable = tqdm(
+                    enumerate(timesteps),
+                    total=len(timesteps),
+                    leave=False,
+                    desc=" " * 4 + "Diffusion denoising",
+                )
+            else:
+                iterable = enumerate(timesteps)
+
+            self.encode_empty_text()
+            text_embed = self.empty_text_embed.repeat((bsz, 1, 1))
+
+            for i, t in iterable:
+                task_emb_anno = torch.tensor([1, 0]).float().unsqueeze(0).to(device)
+                task_emb_anno = torch.cat([torch.sin(task_emb_anno), torch.cos(task_emb_anno)], dim=-1).repeat(bsz, 1)
+
+                if self.scheduler.config.prediction_type == "sample":
+                    feat = self.unet(
+                        rgb_latent, t, text_embed, class_labels=task_emb_anno, return_dict=False, return_decoder_feats=True
+                    )
+                else:
+                    raise
+            
+            feats = (feat,)  # Wrap single tensor in tuple to match expected format for prediction mode
+            feat_for_training = feat  # Keep original for training mode
+        
+        if mode == "predict":
+
+            if isinstance(self.decoder, (nn.parallel.DistributedDataParallel, )):
+
+                preds = self.decoder.module.predict(feats, data_samples, test_cfg=self.test_cfg)
+            
+            else:
+
+                preds = self.decoder.predict(feats, data_samples, test_cfg=self.test_cfg)
+
+            if isinstance(preds, tuple):
+                batch_pred_instances, batch_pred_fields = preds
+            else:
+                batch_pred_instances = preds
+                batch_pred_fields = None
+
+            return self.add_pred_to_datasample(batch_pred_instances=batch_pred_instances, batch_data_samples=data_samples, batch_pred_fields=batch_pred_fields)
+
+        elif mode == "inference":
+            
+            preds = self.decoder.predict(feats, data_samples, test_cfg=self.test_cfg)
+
+            return preds
+            
+        else:
+            if self.test_cfg.get('flip_test', False):
+                
+                if isinstance(feats, list):
+                    # feats is [(feat_orig,), (feat_flipped,)], extract the first one for training
+                    heatmap_pred = self.decoder(feats[0][0])  # Use feat_orig for training
+                else:
+                    # feats is a single tuple (feat,), extract the tensor
+                    heatmap_pred = self.decoder(feats[0])
+            else:
+                # In training mode, decoder expects raw features, not wrapped in tuples
+                heatmap_pred = self.decoder(feat_for_training)  # Use the original feat tensor
+
+            return heatmap_pred
+
+    def encode_empty_text(self):
+
+        """
+        Encode text embedding for empty prompt
+        """
+
+        prompt = ""
+        text_inputs = self.tokenizer(
+            prompt,
+            padding="do_not_pad",
+            max_length=self.tokenizer.model_max_length,
+            truncation=True,
+            return_tensors="pt",
+        )
+        text_input_ids = text_inputs.input_ids.to(self.text_encoder.device)
+        self.empty_text_embed = self.text_encoder(text_input_ids)[0]
+    
+    def encode_rgb(self, rgb_in: torch.Tensor) -> torch.Tensor:
+        """
+        encode RGB image into latent.
+            
+        Args:
+            
+        rgb_in (`torch.Tensor`):
+            Input RGB image to be encoded.
+            
+        Returns:
+            `torch.Tensor`: Image latent.
+        """
+        
+        rgb_latent = self.vae.encode(rgb_in).latent_dist.sample()
+        rgb_latent = rgb_latent * self.rgb_latent_scale_factor
+        
+        return rgb_latent
+
+    def add_pred_to_datasample(self, batch_pred_instances: InstanceList,
+                               batch_pred_fields: Optional[PixelDataList],
+                               batch_data_samples: SampleList) -> SampleList:
+        
+        # from https://github.com/open-mmlab/mmpose/blob/759b39c13fea6ba094afc1fa932f51dc1b11cbf9/mmpose/models/pose_estimators/topdown.py#L122
+        
+        """Add predictions into data samples.
+
+        Args:
+            batch_pred_instances (List[InstanceData]): The predicted instances
+                of the input data batch
+            batch_pred_fields (List[PixelData], optional): The predicted
+                fields (e.g. heatmaps) of the input batch
+            batch_data_samples (List[PoseDataSample]): The input data batch
+
+        Returns:
+            List[PoseDataSample]: A list of data samples where the predictions
+            are stored in the ``pred_instances`` field of each data sample.
+        """
+        assert len(batch_pred_instances) == len(batch_data_samples)
+        if batch_pred_fields is None:
+            batch_pred_fields = []
+        output_keypoint_indices = self.test_cfg.get('output_keypoint_indices',
+                                                    None)
+
+        for pred_instances, pred_fields, data_sample in zip_longest(
+                batch_pred_instances, batch_pred_fields, batch_data_samples):
+
+            gt_instances = data_sample.gt_instances
+
+            bbox_centers = gt_instances.bbox_centers
+            bbox_scales = gt_instances.bbox_scales
+            input_size = data_sample.metainfo['input_size']
+
+            pred_instances.keypoints = pred_instances.keypoints / input_size \
+                * bbox_scales + bbox_centers - 0.5 * bbox_scales
+
+            if output_keypoint_indices is not None:
+                
+                num_keypoints = pred_instances.keypoints.shape[1]
+                for key, value in pred_instances.all_items():
+                    if key.startswith('keypoint'):
+                        pred_instances.set_field(
+                            value[:, output_keypoint_indices], key)
+            
+            pred_instances.bboxes = gt_instances.bboxes
+            pred_instances.bbox_scores = gt_instances.bbox_scores
+
+            data_sample.pred_instances = pred_instances
+
+            if pred_fields is not None:
+                if output_keypoint_indices is not None:
+                    
+                    for key, value in pred_fields.all_items():
+                        if value.shape[0] != num_keypoints:
+                            continue
+                        pred_fields.set_field(value[output_keypoint_indices],
+                                              key)
+                data_sample.pred_fields = pred_fields
+
+        return batch_data_samples

pipelines/__init__.py

@@ -0,0 +1 @@
+# Pipelines package for SDPose OOD

requirements.txt

@@ -8,6 +8,7 @@
 # Core (required)
 opencv-python>=4.8.0
 numpy>=1.24.0
+opencv-python>=4.8.0
 Pillow>=10.0.0
 
 # Optional backend packages (install manually when needed)

scripts/eval.sh

@@ -0,0 +1,40 @@
+#!/bin/bash
+# ============================================
+# Author: T. S. Liang @ Rama Alpaca
+# Emails: tsliang2001@gmail.com | shuangliang@ramaalpaca.com | sliang57@connect.hku.hk
+# Date: Aug. 2025
+# Description: Launch script for evaluating SDPose on pose estimation benchmarks.
+# License: MIT License (see LICENSE for details)
+# ============================================
+
+LAUNCH_EVAL() {
+
+    cd ..
+
+    # ------------------ Configurable Parameters ------------------
+    dataset_name='COCO'
+    keypoint_scheme='body'  # or 'wholebody'
+    dataset_root='/data/coding/datasets'
+    ann_file='/data/coding/datasets/HumanArt/annotations/validation_humanart.json'
+    checkpoint_path='/data/coding/outputs/COCO_813'
+    eval_batch_size=16
+    dataloader_num_workers=16
+    timestep=999
+
+    # ------------------ Runtime Environment ------------------
+    CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 accelerate launch --mixed_precision="bf16"\
+        --multi_gpu \
+        eval/launch_evaluator.py \
+        --dataset_name $dataset_name \
+        --keypoint_scheme $keypoint_scheme \
+        --dataset_root $dataset_root \
+        --checkpoint_path $checkpoint_path \
+        --ann_file $ann_file \
+        --eval_batch_size $eval_batch_size \
+        --dataloader_num_workers $dataloader_num_workers \
+        --enable_xformers_memory_efficient_attention \
+        --timestep $timestep
+}
+
+# ------------- Entry -------------
+LAUNCH_EVAL

utils/EvalMetric.py

@@ -0,0 +1,24 @@
+# Author: T. S. Liang @ Rama Alpaca
+# Emails: tsliang2001@gmail.com | shuangliang@ramaalpaca.com
+# Date: Feb. 2025
+# Description: Build COCO-style keypoint metrics for SDPose evaluation.
+# License: MIT License (see LICENSE for details)
+
+
+from mmpose.evaluation.metrics import CocoMetric, CocoWholeBodyMetric
+
+def get_metric(ann_file, mode = "body"):
+
+    if mode == "body":
+
+        metric = CocoMetric(iou_type='keypoints', ann_file = ann_file)
+
+    elif mode == "wholebody":
+
+        metric = CocoWholeBodyMetric(iou_type='keypoints', ann_file = ann_file)
+
+    else:
+
+        print("Error, please specify the correct type of keypoint scheme.")
+
+    return metric

utils/__init__.py

@@ -0,0 +1 @@
+# Utils package for SDPose OOD

utils/dataset_loader.py

@@ -0,0 +1,179 @@
+# Author: T. S. Liang @ Rama Alpaca
+# Emails: tsliang2001@gmail.com | shuangliang@ramaalpaca.com
+# Date: Feb. 2025
+# Description: The dataloader for the training and evaluation protocol.
+# License: MIT License (see LICENSE for details)
+
+from torch.utils.data import DataLoader
+from mmengine.dataset import default_collate 
+from mmengine.registry import init_default_scope
+from mmpose.registry import DATASETS
+
+init_default_scope('mmpose')
+
+def prepare_dataset(batch_size, dataset_name, dataset_root, mode = "val", num_workers = 4):
+
+    image_size = [768, 1024] ## width x height
+    sigma = 6 ## sigma is 2 for 256
+    scale = 4
+
+    codec = dict(
+        type='UDPHeatmap', input_size=(image_size[0], image_size[1]), heatmap_size=(int(image_size[0]/scale), int(image_size[1]/scale)), sigma=sigma) 
+    
+    train_pipeline = [
+        dict(type='LoadImage'),
+        dict(type='GetBBoxCenterScale'),
+        dict(type='RandomFlip', direction='horizontal'),
+        dict(type='RandomHalfBody'),
+        dict(type='RandomBBoxTransform'),
+        dict(type='TopdownAffine', input_size=codec['input_size'], use_udp=True),
+        dict(type='PhotometricDistortion'),
+        dict(type='CacheRGBTarget'),
+        dict(
+        type='Albumentation',
+        transforms=[
+            dict(type='Blur', p=0.1),
+            dict(type='MedianBlur', p=0.1),
+            dict(
+                type='CoarseDropout',
+                max_holes=1,
+                max_height=0.4,
+                max_width=0.4,
+                min_holes=1,
+                min_height=0.2,
+                min_width=0.2,
+                p=1.0),
+        ]),
+        dict(type='GenerateTarget', encoder=codec),
+        dict(type='PackPoseInputs'),
+    ]
+
+    val_pipeline = [
+        dict(type='LoadImage'),
+        dict(type='GetBBoxCenterScale'),
+        dict(type='TopdownAffine', input_size=codec['input_size'], use_udp=True),
+        dict(type='CacheRGBTarget'),
+        dict(type='PackPoseInputs'),
+    ]
+
+    train_cfg = dict(
+        type='CocoDataset',
+        data_root=dataset_root,
+        data_mode='topdown',
+        ann_file='COCO/annotations/person_keypoints_train2017.json',
+        data_prefix=dict(img='COCO/train2017/'),
+        pipeline=train_pipeline,
+    )
+
+    if dataset_name == "COCO_OOD":
+
+        print("Running Validation on COCO_OOD.")
+
+        val_cfg = dict(
+            type='CocoDataset',
+            data_root=dataset_root,
+            data_mode='topdown',
+            ann_file='COCO/annotations/person_keypoints_val2017.json',
+            data_prefix=dict(img='COCO/val2017oil/'),
+            test_mode=True,
+            pipeline=val_pipeline,
+            bbox_file='COCO/person_detection_results/COCO_val2017_detections_AP_H_70_person.json',
+        )
+    
+    elif dataset_name == "Humanart":
+
+        print("Running Validation on Humanart.")
+
+        val_cfg = dict(
+            type='HumanArtDataset',
+            data_root=dataset_root,
+            data_mode='topdown',
+            ann_file='HumanArt/annotations/validation_humanart.json',
+            test_mode=True,
+            pipeline=val_pipeline,
+        )
+
+    elif dataset_name == "COCO":
+
+        print("Running Validation on COCO.")
+
+        val_cfg = dict(
+            type='CocoDataset',
+            data_root=dataset_root,
+            data_mode='topdown',
+            ann_file='COCO/annotations/person_keypoints_val2017.json',
+            data_prefix=dict(img='COCO/val2017/'),
+            test_mode=True,
+            pipeline=val_pipeline,
+            bbox_file='COCO/person_detection_results/COCO_val2017_detections_AP_H_70_person.json',
+        )
+    
+    elif dataset_name == "COCOWholebody":
+
+        print("Running Validation on COCO Wholebody.")
+
+        val_cfg = dict(
+            type='CocoWholeBodyDataset',
+            data_root=dataset_root,
+            data_mode='topdown',
+            ann_file='COCO/annotations/coco_wholebody_val_v1.0.json',
+            data_prefix=dict(img='COCO/val2017/'),
+            test_mode=True,
+            pipeline=val_pipeline,
+            bbox_file='COCO/person_detection_results/COCO_val2017_detections_AP_H_70_person.json',
+        )
+
+    elif dataset_name == "COCO-OOD_Wholebody":
+
+        print("Running Validation on COCO-OOD Wholebody.")
+
+        val_cfg = dict(
+            type='CocoWholeBodyDataset',
+            data_root=dataset_root,
+            data_mode='topdown',
+            ann_file='COCO/annotations/coco_wholebody_val_v1.0.json',
+            data_prefix=dict(img='COCO/val2017oil/'),
+            test_mode=True,
+            pipeline=val_pipeline,
+            bbox_file='COCO/person_detection_results/COCO_val2017_detections_AP_H_70_person.json',
+        )
+
+    if mode == "train":
+
+        train_ds = DATASETS.build(train_cfg)
+        val_ds = DATASETS.build(val_cfg)
+
+        train_loader = DataLoader(
+            train_ds,
+            batch_size=batch_size,
+            shuffle=True,
+            num_workers=num_workers,
+            collate_fn=default_collate,
+            pin_memory=True,
+        )
+        
+        val_loader = DataLoader(
+            val_ds,
+            batch_size=batch_size,
+            shuffle=False,
+            num_workers=num_workers,
+            collate_fn=default_collate,
+            pin_memory=True,
+        )
+
+        return train_loader, val_loader
+
+    elif mode == "val":
+
+        val_ds = DATASETS.build(val_cfg)
+
+        val_loader = DataLoader(
+            val_ds,
+            batch_size=batch_size,
+            shuffle=False,
+            num_workers=num_workers,
+            collate_fn=default_collate,
+            pin_memory=True,
+        )
+
+        return val_loader