update sample tools

src/sample_utils/__init__.py

@@ -0,0 +1,3 @@
+from .pose_sampler import PoseSampler
+from .beta_sampler import sample_beta
+from .trans_sampler import sample_transl4pp, sample_transl4t, sample_transl4m

src/sample_utils/beta_sampler.py

@@ -0,0 +1,36 @@
+import numpy as np
+import torch
+
+def sample_beta(batch_size=1, sampling_method='normal', range_limit=3.0, device='cpu'):
+    """
+    对SMPL模型的beta参数进行采样，返回1x10的PyTorch张量。
+
+    参数:
+        batch_size (int): 采样样本数量，默认为1。
+        sampling_method (str): 采样方法，'uniform'（均匀采样）或'normal'（正态分布采样）。默认为'normal'。
+        range_limit (float): beta参数的范围限制，默认为3.0（即[-3, 3]）。
+        device (torch.device): 计算设备 (e.g., 'cuda' or 'cpu')
+
+    返回:
+        torch.Tensor: 形状为(batch_size, 10)的beta参数张量。
+
+    异常:
+        ValueError: 如果sampling_method不是'uniform'或'normal'。
+    """
+    beta_dim = 10  # SMPL beta参数维度
+    
+    if sampling_method == 'uniform':
+        # 均匀分布采样
+        beta = np.random.uniform(low=-range_limit, high=range_limit, size=(batch_size, beta_dim))
+    elif sampling_method == 'normal':
+        # 正态分布采样
+        beta = np.random.normal(loc=0, scale=2, size=(batch_size, beta_dim))
+        beta = np.clip(beta, -range_limit, range_limit)  # 限制在[-range_limit, range_limit]
+    else:
+        raise ValueError("sampling_method must be 'uniform' or 'normal'")
+    
+    # 转换为PyTorch张量
+    beta_tensor = torch.tensor(beta, dtype=torch.float32).to(device)
+    
+    return beta_tensor
+

src/sample_utils/lib/ckpt/fm_20251117_172225/checkpoints/best.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4bafdbcfb572a80b3cae8d2df794a9dbb5b5cd4bf2f671eefd1c0549dcdd8ead
+size 59298341

src/sample_utils/lib/ckpt/fm_20251117_172225/config.json

@@ -0,0 +1,26 @@
+{
+    "model_type": "flow_matching",
+    "device": "cuda:0",
+    "epochs": 50,
+    "batch_size": 64,
+    "learning_rate": 0.0001,
+    "weight_decay": 0.0,
+    "grad_clip": 1.0,
+    "num_workers": 0,
+    "timesteps": 1000,
+    "flow_sigma": 0.01,
+    "base_channels": 128,
+    "channel_multipliers": [
+        1,
+        2,
+        4
+    ],
+    "time_emb_dim": 128,
+    "mid_structure": "attention",
+    "mid_num_heads": 4,
+    "output_dir": "outputs",
+    "exp_name": "",
+    "log_interval": 100,
+    "save_interval": 999,
+    "resume": ""
+}

src/sample_utils/lib/ckpt/fm_20251117_172225/logs/training.log

@@ -0,0 +1,170 @@
+[2025-11-17 17:22:25] [INFO] Using device: cuda:0
+[2025-11-17 17:22:25] [INFO] Experiment: 20251117_172225
+[2025-11-17 17:22:25] [INFO] Output directory: outputs/20251117_172225
+[2025-11-17 17:22:25] [INFO] Configuration saved to: outputs/20251117_172225/config.json
+[2025-11-17 17:22:25] [INFO] Loading datasets...
+[2025-11-17 17:23:12] [INFO] Loaded train data: 184000 samples
+[2025-11-17 17:23:12] [INFO] Initializing FLOW_MATCHING model...
+[2025-11-17 17:23:12] [INFO] UNet mid structure: attention
+[2025-11-17 17:23:12] [INFO] Attention heads: 4
+[2025-11-17 17:23:12] [INFO] FlowMatching initialized with sigma=0.01
+[2025-11-17 17:23:12] [INFO] Total parameters: 4,931,329
+[2025-11-17 17:23:12] [INFO] Trainable parameters: 4,931,329
+[2025-11-17 17:23:12] [INFO] Hyperparameters:
+[2025-11-17 17:23:12] [INFO]   model_type: flow_matching
+[2025-11-17 17:23:12] [INFO]   device: cuda:0
+[2025-11-17 17:23:12] [INFO]   epochs: 50
+[2025-11-17 17:23:12] [INFO]   batch_size: 64
+[2025-11-17 17:23:12] [INFO]   learning_rate: 0.0001
+[2025-11-17 17:23:12] [INFO]   weight_decay: 0.0
+[2025-11-17 17:23:12] [INFO]   grad_clip: 1.0
+[2025-11-17 17:23:12] [INFO]   num_workers: 0
+[2025-11-17 17:23:12] [INFO]   timesteps: 1000
+[2025-11-17 17:23:12] [INFO]   flow_sigma: 0.01
+[2025-11-17 17:23:12] [INFO]   base_channels: 128
+[2025-11-17 17:23:12] [INFO]   channel_multipliers: [1, 2, 4]
+[2025-11-17 17:23:12] [INFO]   time_emb_dim: 128
+[2025-11-17 17:23:12] [INFO]   mid_structure: attention
+[2025-11-17 17:23:12] [INFO]   mid_num_heads: 4
+[2025-11-17 17:23:12] [INFO]   output_dir: outputs
+[2025-11-17 17:23:12] [INFO]   exp_name: 
+[2025-11-17 17:23:12] [INFO]   log_interval: 100
+[2025-11-17 17:23:12] [INFO]   save_interval: 999
+[2025-11-17 17:23:12] [INFO]   resume: 
+[2025-11-17 17:23:12] [INFO] Starting training...
+[2025-11-17 17:29:55] [INFO] Epoch [1/50] Average Loss: 1.047125
+[2025-11-17 17:29:55] [INFO] New best model! Loss: 1.047125
+[2025-11-17 17:29:56] [INFO] Checkpoint saved at epoch 1
+[2025-11-17 17:36:55] [INFO] Epoch [2/50] Average Loss: 0.947381
+[2025-11-17 17:36:55] [INFO] New best model! Loss: 0.947381
+[2025-11-17 17:37:10] [INFO] Checkpoint saved at epoch 2
+[2025-11-17 17:44:06] [INFO] Epoch [3/50] Average Loss: 0.922550
+[2025-11-17 17:44:06] [INFO] New best model! Loss: 0.922550
+[2025-11-17 17:44:11] [INFO] Checkpoint saved at epoch 3
+[2025-11-17 17:50:46] [INFO] Epoch [4/50] Average Loss: 0.909522
+[2025-11-17 17:50:46] [INFO] New best model! Loss: 0.909522
+[2025-11-17 17:50:47] [INFO] Checkpoint saved at epoch 4
+[2025-11-17 17:57:23] [INFO] Epoch [5/50] Average Loss: 0.899388
+[2025-11-17 17:57:23] [INFO] New best model! Loss: 0.899388
+[2025-11-17 17:57:24] [INFO] Checkpoint saved at epoch 5
+[2025-11-17 18:04:04] [INFO] Epoch [6/50] Average Loss: 0.893692
+[2025-11-17 18:04:04] [INFO] New best model! Loss: 0.893692
+[2025-11-17 18:04:05] [INFO] Checkpoint saved at epoch 6
+[2025-11-17 18:10:46] [INFO] Epoch [7/50] Average Loss: 0.886531
+[2025-11-17 18:10:46] [INFO] New best model! Loss: 0.886531
+[2025-11-17 18:10:47] [INFO] Checkpoint saved at epoch 7
+[2025-11-17 18:17:26] [INFO] Epoch [8/50] Average Loss: 0.881362
+[2025-11-17 18:17:26] [INFO] New best model! Loss: 0.881362
+[2025-11-17 18:17:27] [INFO] Checkpoint saved at epoch 8
+[2025-11-17 18:24:08] [INFO] Epoch [9/50] Average Loss: 0.878617
+[2025-11-17 18:24:08] [INFO] New best model! Loss: 0.878617
+[2025-11-17 18:24:09] [INFO] Checkpoint saved at epoch 9
+[2025-11-17 18:30:53] [INFO] Epoch [10/50] Average Loss: 0.874885
+[2025-11-17 18:30:53] [INFO] New best model! Loss: 0.874885
+[2025-11-17 18:30:53] [INFO] Checkpoint saved at epoch 10
+[2025-11-17 18:37:29] [INFO] Epoch [11/50] Average Loss: 0.870838
+[2025-11-17 18:37:29] [INFO] New best model! Loss: 0.870838
+[2025-11-17 18:37:30] [INFO] Checkpoint saved at epoch 11
+[2025-11-17 18:44:02] [INFO] Epoch [12/50] Average Loss: 0.868847
+[2025-11-17 18:44:02] [INFO] New best model! Loss: 0.868847
+[2025-11-17 18:44:03] [INFO] Checkpoint saved at epoch 12
+[2025-11-17 18:50:45] [INFO] Epoch [13/50] Average Loss: 0.867175
+[2025-11-17 18:50:45] [INFO] New best model! Loss: 0.867175
+[2025-11-17 18:50:46] [INFO] Checkpoint saved at epoch 13
+[2025-11-17 18:57:22] [INFO] Epoch [14/50] Average Loss: 0.864302
+[2025-11-17 18:57:22] [INFO] New best model! Loss: 0.864302
+[2025-11-17 18:57:23] [INFO] Checkpoint saved at epoch 14
+[2025-11-17 19:04:01] [INFO] Epoch [15/50] Average Loss: 0.862865
+[2025-11-17 19:04:01] [INFO] New best model! Loss: 0.862865
+[2025-11-17 19:04:01] [INFO] Checkpoint saved at epoch 15
+[2025-11-17 19:10:41] [INFO] Epoch [16/50] Average Loss: 0.859902
+[2025-11-17 19:10:41] [INFO] New best model! Loss: 0.859902
+[2025-11-17 19:10:42] [INFO] Checkpoint saved at epoch 16
+[2025-11-17 19:17:17] [INFO] Epoch [17/50] Average Loss: 0.857788
+[2025-11-17 19:17:17] [INFO] New best model! Loss: 0.857788
+[2025-11-17 19:17:18] [INFO] Checkpoint saved at epoch 17
+[2025-11-17 19:24:00] [INFO] Epoch [18/50] Average Loss: 0.855244
+[2025-11-17 19:24:00] [INFO] New best model! Loss: 0.855244
+[2025-11-17 19:24:00] [INFO] Checkpoint saved at epoch 18
+[2025-11-17 19:30:42] [INFO] Epoch [19/50] Average Loss: 0.854479
+[2025-11-17 19:30:42] [INFO] New best model! Loss: 0.854479
+[2025-11-17 19:30:43] [INFO] Checkpoint saved at epoch 19
+[2025-11-17 19:37:24] [INFO] Epoch [20/50] Average Loss: 0.853335
+[2025-11-17 19:37:24] [INFO] New best model! Loss: 0.853335
+[2025-11-17 19:37:25] [INFO] Checkpoint saved at epoch 20
+[2025-11-17 19:44:00] [INFO] Epoch [21/50] Average Loss: 0.851748
+[2025-11-17 19:44:00] [INFO] New best model! Loss: 0.851748
+[2025-11-17 19:44:01] [INFO] Checkpoint saved at epoch 21
+[2025-11-17 19:50:42] [INFO] Epoch [22/50] Average Loss: 0.849092
+[2025-11-17 19:50:42] [INFO] New best model! Loss: 0.849092
+[2025-11-17 19:50:42] [INFO] Checkpoint saved at epoch 22
+[2025-11-17 19:57:25] [INFO] Epoch [23/50] Average Loss: 0.847517
+[2025-11-17 19:57:25] [INFO] New best model! Loss: 0.847517
+[2025-11-17 19:57:25] [INFO] Checkpoint saved at epoch 23
+[2025-11-17 20:04:06] [INFO] Epoch [24/50] Average Loss: 0.845896
+[2025-11-17 20:04:06] [INFO] New best model! Loss: 0.845896
+[2025-11-17 20:04:07] [INFO] Checkpoint saved at epoch 24
+[2025-11-17 20:10:44] [INFO] Epoch [25/50] Average Loss: 0.846312
+[2025-11-17 20:17:26] [INFO] Epoch [26/50] Average Loss: 0.844638
+[2025-11-17 20:17:26] [INFO] New best model! Loss: 0.844638
+[2025-11-17 20:17:27] [INFO] Checkpoint saved at epoch 26
+[2025-11-17 20:24:07] [INFO] Epoch [27/50] Average Loss: 0.842196
+[2025-11-17 20:24:07] [INFO] New best model! Loss: 0.842196
+[2025-11-17 20:24:08] [INFO] Checkpoint saved at epoch 27
+[2025-11-17 20:30:50] [INFO] Epoch [28/50] Average Loss: 0.840379
+[2025-11-17 20:30:50] [INFO] New best model! Loss: 0.840379
+[2025-11-17 20:30:51] [INFO] Checkpoint saved at epoch 28
+[2025-11-17 20:37:33] [INFO] Epoch [29/50] Average Loss: 0.840037
+[2025-11-17 20:37:33] [INFO] New best model! Loss: 0.840037
+[2025-11-17 20:37:34] [INFO] Checkpoint saved at epoch 29
+[2025-11-17 20:44:17] [INFO] Epoch [30/50] Average Loss: 0.839178
+[2025-11-17 20:44:17] [INFO] New best model! Loss: 0.839178
+[2025-11-17 20:44:18] [INFO] Checkpoint saved at epoch 30
+[2025-11-17 20:50:56] [INFO] Epoch [31/50] Average Loss: 0.835445
+[2025-11-17 20:50:56] [INFO] New best model! Loss: 0.835445
+[2025-11-17 20:50:57] [INFO] Checkpoint saved at epoch 31
+[2025-11-17 20:57:37] [INFO] Epoch [32/50] Average Loss: 0.837469
+[2025-11-17 21:04:20] [INFO] Epoch [33/50] Average Loss: 0.835059
+[2025-11-17 21:04:20] [INFO] New best model! Loss: 0.835059
+[2025-11-17 21:04:21] [INFO] Checkpoint saved at epoch 33
+[2025-11-17 21:11:02] [INFO] Epoch [34/50] Average Loss: 0.835552
+[2025-11-17 21:17:46] [INFO] Epoch [35/50] Average Loss: 0.832502
+[2025-11-17 21:17:46] [INFO] New best model! Loss: 0.832502
+[2025-11-17 21:17:47] [INFO] Checkpoint saved at epoch 35
+[2025-11-17 21:24:42] [INFO] Epoch [36/50] Average Loss: 0.833365
+[2025-11-17 21:31:20] [INFO] Epoch [37/50] Average Loss: 0.831459
+[2025-11-17 21:31:20] [INFO] New best model! Loss: 0.831459
+[2025-11-17 21:31:21] [INFO] Checkpoint saved at epoch 37
+[2025-11-17 21:37:55] [INFO] Epoch [38/50] Average Loss: 0.829138
+[2025-11-17 21:37:55] [INFO] New best model! Loss: 0.829138
+[2025-11-17 21:37:56] [INFO] Checkpoint saved at epoch 38
+[2025-11-17 21:44:30] [INFO] Epoch [39/50] Average Loss: 0.829815
+[2025-11-17 21:51:10] [INFO] Epoch [40/50] Average Loss: 0.829309
+[2025-11-17 21:57:46] [INFO] Epoch [41/50] Average Loss: 0.828454
+[2025-11-17 21:57:46] [INFO] New best model! Loss: 0.828454
+[2025-11-17 21:57:47] [INFO] Checkpoint saved at epoch 41
+[2025-11-17 22:04:26] [INFO] Epoch [42/50] Average Loss: 0.828211
+[2025-11-17 22:04:26] [INFO] New best model! Loss: 0.828211
+[2025-11-17 22:04:27] [INFO] Checkpoint saved at epoch 42
+[2025-11-17 22:11:06] [INFO] Epoch [43/50] Average Loss: 0.828369
+[2025-11-17 22:17:38] [INFO] Epoch [44/50] Average Loss: 0.827580
+[2025-11-17 22:17:38] [INFO] New best model! Loss: 0.827580
+[2025-11-17 22:18:00] [INFO] Checkpoint saved at epoch 44
+[2025-11-17 22:24:56] [INFO] Epoch [45/50] Average Loss: 0.826730
+[2025-11-17 22:24:56] [INFO] New best model! Loss: 0.826730
+[2025-11-17 22:24:57] [INFO] Checkpoint saved at epoch 45
+[2025-11-17 22:32:10] [INFO] Epoch [46/50] Average Loss: 0.826514
+[2025-11-17 22:32:10] [INFO] New best model! Loss: 0.826514
+[2025-11-17 22:32:10] [INFO] Checkpoint saved at epoch 46
+[2025-11-17 22:39:14] [INFO] Epoch [47/50] Average Loss: 0.826701
+[2025-11-17 22:45:49] [INFO] Epoch [48/50] Average Loss: 0.825706
+[2025-11-17 22:45:49] [INFO] New best model! Loss: 0.825706
+[2025-11-17 22:45:50] [INFO] Checkpoint saved at epoch 48
+[2025-11-17 22:52:25] [INFO] Epoch [49/50] Average Loss: 0.824755
+[2025-11-17 22:52:25] [INFO] New best model! Loss: 0.824755
+[2025-11-17 22:52:26] [INFO] Checkpoint saved at epoch 49
+[2025-11-17 22:59:00] [INFO] Epoch [50/50] Average Loss: 0.827608
+[2025-11-17 22:59:01] [INFO] Checkpoint saved at epoch 50
+[2025-11-17 22:59:01] [INFO] Training completed!
+[2025-11-17 22:59:01] [INFO] Best loss: 0.824755
+[2025-11-17 22:59:01] [INFO] Checkpoints saved in: outputs/20251117_172225/checkpoints

src/sample_utils/lib/data_stats/pose_stats.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:99993ef6d0e6359054c8dd3866f9a63951b85d4e6c5b5999a98f11fe01246bf5
+size 1962

src/sample_utils/lib/data_stats/pp_pose_stats.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:62c10d637ce1543d895dbc6ccee898b1f44e9a0e0672f26432f857266d853875
+size 1980

src/sample_utils/lib/data_stats/t_pose_stats.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8b5264858353aadcebf6547da84dcb3a3c046682ee1036fe047a78f983449fb1
+size 1974

src/sample_utils/lib/model/ddpm.py

@@ -0,0 +1,194 @@
+# motion_generation/lib/model/ddpm.py
+
+import torch
+import torch.nn as nn
+from lib.model.unet1d import UNet1D
+
+
+class DDPM(nn.Module):
+    def __init__(self, model: UNet1D, beta_schedule='linear', timesteps=1000):
+        super().__init__()
+        self.model = model  # 噪声预测器 (UNet1D)
+        self.timesteps = timesteps
+
+        # ----------------------------------------------------
+        # 1. 定义 Beta 调度和相关参数
+        # ----------------------------------------------------
+        
+        # 1.1. 计算 Betas (噪声水平)
+        if beta_schedule == 'linear':
+            betas = self._linear_beta_schedule(timesteps)
+        else:
+            raise NotImplementedError(f"不支持的 beta_schedule: {beta_schedule}")
+
+        # 将所有参数注册为 buffer (不会被训练，但会保存到 state_dict)
+        self.register_buffer('betas', betas)
+        
+        # 1.2. 计算 Alpha 相关的参数
+        alphas = 1.0 - betas
+        alphas_cumprod = torch.cumprod(alphas, dim=0) # 累积乘积: \bar{\alpha}_t
+        alphas_cumprod_prev = torch.cat([torch.tensor([1.0]), alphas_cumprod[:-1]])
+
+        # 方便计算的根号形式
+        sqrt_alphas_cumprod = torch.sqrt(alphas_cumprod)
+        sqrt_one_minus_alphas_cumprod = torch.sqrt(1.0 - alphas_cumprod)
+        sqrt_recip_alphas = torch.sqrt(1.0 / alphas)  # 1/sqrt(alpha_t)
+
+        # 后验均值计算所需的参数 (反向过程)
+        posterior_variance = betas * (1.0 - alphas_cumprod_prev) / (1.0 - alphas_cumprod)
+
+        # 注册 buffers
+        self.register_buffer('alphas', alphas)
+        self.register_buffer('alphas_cumprod', alphas_cumprod)
+        self.register_buffer('sqrt_alphas_cumprod', sqrt_alphas_cumprod)
+        self.register_buffer('sqrt_one_minus_alphas_cumprod', sqrt_one_minus_alphas_cumprod)
+        self.register_buffer('sqrt_recip_alphas', sqrt_recip_alphas)
+        self.register_buffer('posterior_variance', posterior_variance)
+        
+        # ----------------------------------------------------
+
+    def _linear_beta_schedule(self, timesteps, start=0.0001, end=0.02):
+        """线性 Beta 调度"""
+        return torch.linspace(start, end, timesteps)
+
+    def get_index_from_t(self, variances, t, x_shape):
+        """从 (T,) 维度的张量中，根据时间步 t (B,) 提取对应的方差/系数，并重塑至 (B, 1, 1...)"""
+        B = t.shape[0]
+        out = variances.gather(-1, t) # 从 variances 中提取 t 对应的值
+
+        # 重塑维度以匹配 x 的形状 (B, C) -> (B, C)
+        # 对于我们的姿态数据 (B, 72)， C=72，所以只需要重塑到 (B, 1)
+        return out.reshape(B, *([1] * (len(x_shape) - 1)))
+
+
+    # ----------------------------------------------------
+    # 2. 前向扩散 (加噪)
+    # ----------------------------------------------------
+    def forward_diffusion(self, x_start, t, noise=None):
+        """
+        前向过程：计算 x_t = sqrt(alpha_bar_t) * x_0 + sqrt(1 - alpha_bar_t) * noise
+        
+        Args:
+            x_start: 初始数据 x_0 (B, 72)
+            t: 时间步 (B,)
+            noise: 用于加噪的噪声，如果为 None 则随机生成
+        
+        Returns:
+            x_t: t时刻加噪后的数据 (B, 72)
+            noise: 实际使用的噪声 (B, 72)
+        """
+        if noise is None:
+            noise = torch.randn_like(x_start)
+            
+        # 提取 sqrt(alpha_bar_t) 和 sqrt(1 - alpha_bar_t)
+        sqrt_alpha_bar_t = self.get_index_from_t(self.sqrt_alphas_cumprod, t, x_start.shape)
+        sqrt_one_minus_alpha_bar_t = self.get_index_from_t(self.sqrt_one_minus_alphas_cumprod, t, x_start.shape)
+        
+        # 公式实现
+        x_t = sqrt_alpha_bar_t * x_start + sqrt_one_minus_alpha_bar_t * noise
+        return x_t, noise
+
+    # ----------------------------------------------------
+    # 3. 损失计算 (训练过程)
+    # ----------------------------------------------------
+    def forward(self, x_start):
+        """
+        训练时的前向传播，用于计算损失。
+        
+        Args:
+            x_start: 批量原始数据 (归一化后的姿态) (B, 72)
+        
+        Returns:
+            loss: 均方误差 (MSE) 损失
+        """
+        B = x_start.shape[0]
+        
+        # 1. 随机采样时间步 t
+        t = torch.randint(0, self.timesteps, (B,), device=x_start.device).long()
+        
+        # 2. 随机生成噪声
+        noise = torch.randn_like(x_start)
+        
+        # 3. 计算 x_t 和 t 时刻的真实噪声
+        x_t, _ = self.forward_diffusion(x_start, t, noise)
+        
+        # 4. 噪声预测网络预测噪声
+        predicted_noise = self.model(x_t, t)
+        
+        # 5. 计算损失：预测噪声和真实噪声的均方误差 (MSE)
+        loss = nn.functional.mse_loss(predicted_noise, noise)
+        
+        return loss
+
+    # ----------------------------------------------------
+    # 4. 采样 (生成过程)
+    # ----------------------------------------------------
+    @torch.no_grad()
+    def sample(self, sample_shape, device, return_intermediates=False, save_interval=None):
+        """
+        逆向过程：从纯噪声开始，迭代去噪 T 步，生成新的姿态数据。
+
+        Args:
+            sample_shape: 要生成的样本形状 (B, 72)
+            device: 运行设备
+            return_intermediates: 是否返回中间状态
+            save_interval: 保存中间状态的间隔（timesteps），如果为None则不保存
+
+        Returns:
+            如果 return_intermediates=False:
+                x_0: 最终生成的归一化姿态 (B, 72)
+            如果 return_intermediates=True:
+                (x_0, intermediates): 最终姿态和中间状态列表
+                intermediates: List[(timestep, x_t)] 包含时间步和对应的状态
+        """
+        # 从纯噪声开始 x_T ~ N(0, I)
+        x = torch.randn(sample_shape, device=device)
+
+        # 存储中间状态
+        intermediates = []
+        if return_intermediates:
+            # 保存初始噪声状态
+            intermediates.append((self.timesteps, x.clone().cpu()))
+
+        # 从 T-1 步迭代到 0 步
+        for t in reversed(range(0, self.timesteps)):
+            t_tensor = torch.full((sample_shape[0],), t, device=device, dtype=torch.long)
+
+            # 1. 预测噪声
+            predicted_noise = self.model(x, t_tensor)
+
+            # 2. 提取当前时刻的系数
+            beta_t = self.get_index_from_t(self.betas, t_tensor, x.shape)
+            sqrt_one_minus_alpha_bar_t = self.get_index_from_t(self.sqrt_one_minus_alphas_cumprod, t_tensor, x.shape)
+            sqrt_recip_alpha_t = self.get_index_from_t(self.sqrt_recip_alphas, t_tensor, x.shape)
+
+            # 3. 计算均值 (mu_t-1)
+            # DDPM公式: μ_t = (1/sqrt(alpha_t)) * (x_t - (beta_t/sqrt(1-alpha_bar_t)) * predicted_noise)
+            mean = sqrt_recip_alpha_t * (x - beta_t * predicted_noise / sqrt_one_minus_alpha_bar_t)
+
+            # 4. 添加噪声项
+            if t > 0:
+                variance = self.get_index_from_t(self.posterior_variance, t_tensor, x.shape)
+                noise = torch.randn_like(x)
+                # x_{t-1} = \mu_{t-1} + \sigma_{t-1} * z
+                x = mean + torch.sqrt(variance) * noise
+            else:
+                x = mean # t=0 时不加噪声，直接取均值作为最终输出
+
+            # 保存中间状态
+            if return_intermediates:
+                if save_interval is None:
+                    # 保存所有状态
+                    intermediates.append((t, x.clone().cpu()))
+                elif t % save_interval == 0 or t == 0:
+                    # 按间隔保存
+                    intermediates.append((t, x.clone().cpu()))
+
+        # 将输出限制在 [-1, 1] 附近（可选，取决于您的归一化范围）
+        # x = x.clamp(-1., 1.)
+
+        if return_intermediates:
+            return x, intermediates
+        return x
+    
+    

src/sample_utils/lib/model/flow_matching.py

@@ -0,0 +1,248 @@
+"""
+Flow Matching implementation for continuous normalizing flows.
+
+Flow matching learns a velocity field v(x, t) that transports samples from a noise
+distribution to the data distribution via ordinary differential equations (ODEs).
+
+Reference: "Flow Matching for Generative Modeling" (Lipman et al., 2023)
+"""
+
+import torch
+import torch.nn as nn
+from tqdm import tqdm
+
+
+class FlowMatching(nn.Module):
+    """
+    Flow Matching model that learns to transport noise to data via velocity fields.
+
+    Args:
+        model: Neural network that predicts velocity field v(x, t)
+               Should take (x, t) and return predicted velocity of same shape as x
+        sigma: Standard deviation for conditional flow matching (default: 0.0)
+               When sigma > 0, uses conditional flow matching with Gaussian paths
+    """
+
+    def __init__(self, model, sigma=0.0):
+        super().__init__()
+        self.model = model
+        self.sigma = sigma
+
+    def forward(self, x_0, return_loss=True):
+        """
+        Compute flow matching loss for a batch of data.
+
+        Forward process:
+            - Sample t uniformly from [0, 1]
+            - Sample noise x_1 ~ N(0, I)
+            - Interpolate: x_t = t * x_0 + (1 - t) * x_1
+            - True velocity: v_t = x_0 - x_1
+            - Loss: MSE(predicted_velocity, true_velocity)
+
+        Args:
+            x_0: Clean data samples (B, D)
+            return_loss: If True, return scalar loss. If False, return per-sample losses
+
+        Returns:
+            loss: Scalar loss if return_loss=True, else (B,) tensor of per-sample losses
+        """
+        batch_size = x_0.shape[0]
+        device = x_0.device
+
+        # Sample random timesteps uniformly from [0, 1]
+        t = torch.rand(batch_size, device=device)
+
+        # Sample noise from standard normal
+        x_1 = torch.randn_like(x_0)
+
+        # Interpolate between noise and data
+        # x_t = t * x_0 + (1 - t) * x_1
+        t_expanded = t.view(batch_size, *([1] * (x_0.ndim - 1)))  # (B, 1, 1, ...)
+        x_t = t_expanded * x_0 + (1 - t_expanded) * x_1
+
+        # True velocity field: dx_t/dt = x_0 - x_1
+        # This is the derivative of the linear interpolation
+        true_velocity = x_0 - x_1
+
+        # Add Gaussian noise for conditional flow matching (if sigma > 0)
+        if self.sigma > 0:
+            x_t = x_t + self.sigma * torch.randn_like(x_t)
+
+        # Predict velocity using the model
+        predicted_velocity = self.model(x_t, t)
+
+        # Compute MSE loss
+        loss = (predicted_velocity - true_velocity) ** 2
+
+        if return_loss:
+            return loss.mean()
+        else:
+            # Return per-sample loss (averaged over dimensions)
+            return loss.view(batch_size, -1).mean(dim=1)
+
+    @torch.no_grad()
+    def sample(self,
+               sample_shape,
+               device='cuda',
+               num_steps=100,
+               method='euler',
+               return_intermediates=False,
+               save_interval=10,
+               verbose=True):
+        """
+        Generate samples by solving the ODE: dx/dt = v(x, t) from t=0 to t=1.
+
+        Args:
+            sample_shape: Shape of samples to generate (B, D)
+            device: Device to generate samples on
+            num_steps: Number of integration steps (default: 100)
+            method: ODE solver method ('euler', 'midpoint', 'rk4')
+            return_intermediates: If True, return intermediate states
+            save_interval: Save intermediate states every N steps (if return_intermediates=True)
+            verbose: Show progress bar
+
+        Returns:
+            samples: Generated samples (B, D)
+            intermediates: List of (t, x_t) tuples (if return_intermediates=True)
+        """
+        # Start from noise at t=0
+        x = torch.randn(sample_shape, device=device)
+
+        # Time steps from 0 to 1
+        timesteps = torch.linspace(0, 1, num_steps + 1, device=device)
+        dt = 1.0 / num_steps
+
+        intermediates = []
+        if return_intermediates:
+            intermediates.append((0.0, x.cpu().clone()))
+
+        # Progress bar
+        iterator = tqdm(range(num_steps), desc='Sampling') if verbose else range(num_steps)
+
+        for i in iterator:
+            t = timesteps[i]
+
+            # Create batch of timesteps
+            t_batch = torch.full((sample_shape[0],), t, device=device)
+
+            if method == 'euler':
+                # Euler method: x_{t+dt} = x_t + dt * v(x_t, t)
+                v = self.model(x, t_batch)
+                x = x + dt * v
+
+            elif method == 'midpoint':
+                # Midpoint method (RK2)
+                # k1 = v(x_t, t)
+                # k2 = v(x_t + 0.5*dt*k1, t + 0.5*dt)
+                # x_{t+dt} = x_t + dt * k2
+                k1 = self.model(x, t_batch)
+                t_mid = t + 0.5 * dt
+                t_mid_batch = torch.full((sample_shape[0],), t_mid, device=device)
+                k2 = self.model(x + 0.5 * dt * k1, t_mid_batch)
+                x = x + dt * k2
+
+            elif method == 'rk4':
+                # Classic RK4 method
+                t_half = t + 0.5 * dt
+                t_next = t + dt
+
+                t_batch_half = torch.full((sample_shape[0],), t_half, device=device)
+                t_batch_next = torch.full((sample_shape[0],), t_next, device=device)
+
+                k1 = self.model(x, t_batch)
+                k2 = self.model(x + 0.5 * dt * k1, t_batch_half)
+                k3 = self.model(x + 0.5 * dt * k2, t_batch_half)
+                k4 = self.model(x + dt * k3, t_batch_next)
+
+                x = x + (dt / 6.0) * (k1 + 2*k2 + 2*k3 + k4)
+
+            else:
+                raise ValueError(f"Unknown method: {method}. Choose from 'euler', 'midpoint', 'rk4'")
+
+            # Save intermediate states
+            if return_intermediates and (i + 1) % save_interval == 0:
+                intermediates.append((timesteps[i + 1].item(), x.cpu().clone()))
+
+        # Final state at t=1
+        if return_intermediates:
+            if len(intermediates) == 0 or intermediates[-1][0] != 1.0:
+                intermediates.append((1.0, x.cpu().clone()))
+            return x, intermediates
+
+        return x
+
+    @torch.no_grad()
+    def sample_ode(self,
+                   sample_shape,
+                   device='cuda',
+                   rtol=1e-5,
+                   atol=1e-5,
+                   method='dopri5',
+                   return_intermediates=False,
+                   verbose=True):
+        """
+        Generate samples using adaptive ODE solvers from torchdiffeq.
+
+        This method uses scipy-style adaptive solvers that automatically
+        adjust step sizes for accuracy. Requires: pip install torchdiffeq
+
+        Args:
+            sample_shape: Shape of samples to generate (B, D)
+            device: Device to generate samples on
+            rtol: Relative tolerance for ODE solver
+            atol: Absolute tolerance for ODE solver
+            method: Solver method ('dopri5', 'dopri8', 'adams', 'rk4', etc.)
+            return_intermediates: If True, return trajectory
+            verbose: Print status messages
+
+        Returns:
+            samples: Generated samples (B, D)
+            trajectory: Full trajectory if return_intermediates=True
+        """
+        try:
+            from torchdiffeq import odeint
+        except ImportError:
+            raise ImportError(
+                "torchdiffeq is required for adaptive ODE solvers. "
+                "Install with: pip install torchdiffeq"
+            )
+
+        if verbose:
+            print(f"Sampling with adaptive ODE solver: {method}")
+
+        # Start from noise at t=0
+        x_0 = torch.randn(sample_shape, device=device)
+
+        # Define velocity field function
+        def velocity_fn(t, x):
+            # t is a scalar tensor, need to broadcast to batch
+            t_batch = torch.full((sample_shape[0],), t.item(), device=device)
+            return self.model(x, t_batch)
+
+        # Integration time points
+        if return_intermediates:
+            # Return full trajectory with 100 points
+            t_span = torch.linspace(0, 1, 100, device=device)
+        else:
+            # Just start and end points
+            t_span = torch.tensor([0.0, 1.0], device=device)
+
+        # Solve ODE
+        trajectory = odeint(
+            velocity_fn,
+            x_0,
+            t_span,
+            rtol=rtol,
+            atol=atol,
+            method=method
+        )
+
+        # Extract final state
+        x_final = trajectory[-1]
+
+        if return_intermediates:
+            # Convert trajectory to list of (t, x) tuples
+            intermediates = [(t.item(), x.cpu()) for t, x in zip(t_span, trajectory)]
+            return x_final, intermediates
+
+        return x_final

src/sample_utils/lib/model/unet1d.py

@@ -0,0 +1,259 @@
+# motion_generation/lib/model/unet1d.py
+
+import torch
+import torch.nn as nn
+import math
+
+class PositionalEncoding(nn.Module):
+    """
+    用于编码时间步 t 的标准 Transformer 位置编码
+    """
+    def __init__(self, dim):
+        super().__init__()
+        self.dim = dim
+        self.register_buffer('inv_freq', 1. / (10000 ** (torch.arange(0, dim, 2).float() / dim)))
+
+    def forward(self, x):
+        # x 形状: (B,) 时间步索引
+        sinusoid_inp = torch.einsum('i, j -> i j', x.float(), self.inv_freq)
+        emb = torch.cat((sinusoid_inp.sin(), sinusoid_inp.cos()), dim=-1)
+        return emb
+
+class ResBlock1D(nn.Module):
+    """
+    一维残差块，包含 time embedding 的融合
+    """
+    def __init__(self, in_channels, out_channels, time_dim, kernel_size=3):
+        super().__init__()
+        padding = kernel_size // 2
+
+        self.conv1 = nn.Conv1d(in_channels, out_channels, kernel_size, padding=padding)
+        self.bn1 = nn.BatchNorm1d(out_channels)
+        self.act1 = nn.GELU()
+
+        self.conv2 = nn.Conv1d(out_channels, out_channels, kernel_size, padding=padding)
+        self.bn2 = nn.BatchNorm1d(out_channels)
+        self.act2 = nn.GELU()
+
+        # 时间步嵌入层
+        self.time_proj = nn.Linear(time_dim, out_channels)
+
+        # 确保输入/输出通道匹配
+        self.residual_conv = nn.Conv1d(in_channels, out_channels, 1) if in_channels != out_channels else nn.Identity()
+
+    def forward(self, x, time_emb):
+        # x: (B, C_in, L) L=72, C_in=通道数
+
+        h = self.conv1(x)
+        h = self.bn1(h)
+        h = self.act1(h)
+
+        # 融合时间嵌入：沿特征维度广播并相加
+        time_emb_proj = self.time_proj(time_emb).unsqueeze(-1) # (B, C_out) -> (B, C_out, 1)
+        h = h + time_emb_proj
+
+        h = self.conv2(h)
+        h = self.bn2(h)
+        h = self.act2(h)
+
+        return h + self.residual_conv(x)
+
+
+class SelfAttention1D(nn.Module):
+    """
+    一维自注意力模块，包含 time embedding 的融合
+    """
+    def __init__(self, channels, time_dim, num_heads=4):
+        super().__init__()
+        self.channels = channels
+        self.num_heads = num_heads
+        assert channels % num_heads == 0, "channels must be divisible by num_heads"
+
+        # Group normalization for better stability
+        self.norm = nn.GroupNorm(num_groups=8, num_channels=channels)
+
+        # Multi-head attention components
+        self.qkv = nn.Conv1d(channels, channels * 3, 1)
+        self.proj = nn.Conv1d(channels, channels, 1)
+
+        # Time embedding projection
+        self.time_proj = nn.Linear(time_dim, channels)
+
+    def forward(self, x, time_emb):
+        # x: (B, C, L)
+        B, C, L = x.shape
+
+        # Normalize input
+        h = self.norm(x)
+
+        # Add time embedding
+        time_emb_proj = self.time_proj(time_emb).unsqueeze(-1)  # (B, C, 1)
+        h = h + time_emb_proj
+
+        # Compute Q, K, V
+        qkv = self.qkv(h)  # (B, 3*C, L)
+        qkv = qkv.reshape(B, 3, self.num_heads, C // self.num_heads, L)
+        qkv = qkv.permute(1, 0, 2, 4, 3)  # (3, B, num_heads, L, head_dim)
+        q, k, v = qkv[0], qkv[1], qkv[2]
+
+        # Attention
+        scale = (C // self.num_heads) ** -0.5
+        attn = torch.matmul(q, k.transpose(-2, -1)) * scale  # (B, num_heads, L, L)
+        attn = torch.softmax(attn, dim=-1)
+
+        # Apply attention to values
+        out = torch.matmul(attn, v)  # (B, num_heads, L, head_dim)
+        out = out.permute(0, 1, 3, 2).reshape(B, C, L)  # (B, C, L)
+
+        # Project and add residual
+        out = self.proj(out)
+        return out + x
+
+
+class UNet1D(nn.Module):
+    def __init__(self, pose_dim=72, base_channels=128, channel_multipliers=[1, 2, 4],
+                 time_emb_dim=256, mid_structure='conv', mid_num_heads=4):
+        """
+        UNet1D model for 1D sequence processing
+
+        Args:
+            pose_dim: Dimension of input pose (default: 72)
+            base_channels: Base number of channels (default: 128)
+            channel_multipliers: Channel multipliers for each level (default: [1, 2, 4])
+            time_emb_dim: Time embedding dimension (default: 256)
+            mid_structure: Structure for middle layer, either 'conv' or 'attention' (default: 'conv')
+            mid_num_heads: Number of attention heads for mid layer when using attention (default: 4)
+        """
+        super().__init__()
+
+        self.mid_structure = mid_structure
+        assert mid_structure in ['conv', 'attention'], "mid_structure must be 'conv' or 'attention'"
+
+        # 1. Time Embedding
+        self.time_mlp = nn.Sequential(
+            PositionalEncoding(base_channels),
+            nn.Linear(base_channels, time_emb_dim),
+            nn.GELU(),
+            nn.Linear(time_emb_dim, time_emb_dim)
+        )
+
+        # 调整输入通道：姿态向量 (B, 72) -> (B, 1, 72)
+        # 我们将特征维度 (72) 视为长度 L，将 1 视为通道 C
+        in_channels = 1
+
+        # 2. 编码器 (Downsampling)
+        channels = [in_channels] + [base_channels * m for m in channel_multipliers]
+        self.downs = nn.ModuleList()
+
+        for i in range(len(channel_multipliers)):
+            in_c = channels[i]
+            out_c = channels[i+1]
+
+            self.downs.append(nn.ModuleList([
+                ResBlock1D(in_c if i == 0 else in_c, out_c, time_emb_dim), # 输入是 1, L 或 C_in, L
+                nn.MaxPool1d(2) if i < len(channel_multipliers) - 1 else nn.Identity()
+            ]))
+
+        # 3. 中间层 - 根据 mid_structure 选择使用卷积或自注意力
+        mid_c = channels[-1]
+        if mid_structure == 'conv':
+            self.mid = ResBlock1D(mid_c, mid_c, time_emb_dim)
+        elif mid_structure == 'attention':
+            # 确保通道数能被注意力头数整除
+            assert mid_c % mid_num_heads == 0, f"mid_c ({mid_c}) must be divisible by mid_num_heads ({mid_num_heads})"
+            self.mid = SelfAttention1D(mid_c, time_emb_dim, num_heads=mid_num_heads)
+        
+        # 4. 解码器 (Upsampling)
+        self.ups = nn.ModuleList()
+        reversed_channels = list(reversed(channels))
+
+        for i in range(len(channel_multipliers)):
+            in_c = reversed_channels[i]  # 来自上一层的通道数
+            # 最后一层输出 base_channels，而不是 in_channels (1)
+            out_c = reversed_channels[i+1] if i < len(channel_multipliers) - 1 else base_channels
+            skip_c = in_c  # 跳跃连接通道数（来自对应编码器层）
+
+            self.ups.append(nn.ModuleList([
+                # ResBlock 接收拼接后的通道: in_c(来自上层) + skip_c(来自编码器)
+                # 输出为 out_c 通道
+                ResBlock1D(in_c + skip_c, out_c, time_emb_dim),
+                # 上采样到下一层的空间尺寸
+                nn.ConvTranspose1d(out_c, out_c, kernel_size=2, stride=2) if i < len(channel_multipliers) - 1 else nn.Identity(),
+            ]))
+
+        # 5. 输出层 (回到 1 个通道)
+        self.out_conv = nn.Conv1d(base_channels, in_channels, kernel_size=1)
+        
+    def forward(self, x, t):
+        # import pdb; pdb.set_trace()
+        # x: (B, 72) 归一化姿态，t: (B,) 时间步索引
+        x = x.unsqueeze(1) # (B, 1, 72)
+
+        # 1. Time Embedding
+        time_emb = self.time_mlp(t)
+        
+        # 2. 编码器
+        skips = []
+        for resblock, downsample in self.downs:
+            x = resblock(x, time_emb)
+            skips.append(x)
+            x = downsample(x)
+
+        # 3. 中间层
+        x = self.mid(x, time_emb)
+        
+        # 4. 解码器
+        for i, (resblock, upsample) in enumerate(self.ups):
+            skip = skips.pop()
+
+            # 跳跃连接
+            # 检查维度是否匹配，如果 MaxPool 导致了奇数/偶数长度不匹配，需要裁剪
+            if x.shape[-1] != skip.shape[-1]:
+                x = nn.functional.pad(x, (0, skip.shape[-1] - x.shape[-1]))
+
+            # 拼接跳跃连接
+            x = torch.cat((x, skip), dim=1) # 沿通道维度拼接
+
+            # 处理拼接后的特征
+            x = resblock(x, time_emb)
+
+            # 上采样到下一层的空间尺寸（最后一层不上采样）
+            if i < len(self.ups) - 1:
+                x = upsample(x) # (B, C, L) -> (B, C, L*2)
+
+        # 5. 输出层
+        x = self.out_conv(x) # (B, 1, 72)
+        
+        return x.squeeze(1) # (B, 72)
+
+
+if __name__ == "__main__":
+    # 测试 UNet1D 模型
+    print("Testing UNet1D with conv mid structure...")
+    model_conv = UNet1D(mid_structure='conv')
+
+    x = torch.randn(4, 72) # 模拟输入 (B, 72)
+    t = torch.randint(0, 1000, (4,)) # 模拟时间步索引 (B,)
+
+    output_conv = model_conv(x, t)
+    print(f"Input shape: {x.shape}")
+    print(f"Output shape (conv): {output_conv.shape}")
+    print(f"Expected output shape: (4, 72)")
+    print(f"Test passed (conv): {output_conv.shape == torch.Size([4, 72])}")
+
+    print("\nTesting UNet1D with attention mid structure...")
+    model_attn = UNet1D(mid_structure='attention', base_channels=128, mid_num_heads=8)
+
+    output_attn = model_attn(x, t)
+    print(f"Output shape (attention): {output_attn.shape}")
+    print(f"Test passed (attention): {output_attn.shape == torch.Size([4, 72])}")
+
+    # 统计参数量
+    conv_params = sum(p.numel() for p in model_conv.parameters())
+    attn_params = sum(p.numel() for p in model_attn.parameters())
+    print(f"\nParameters (conv): {conv_params:,}")
+    print(f"Parameters (attention): {attn_params:,}")
+    print(f"Difference: {abs(attn_params - conv_params):,}")
+
+
+

src/sample_utils/pose_sampler.py

@@ -0,0 +1,138 @@
+import torch
+import os
+from pathlib import Path
+import numpy as np
+import json
+import sys
+
+# PROJECT_ROOT = Path(__file__).resolve().parent.parent
+# sys.path.insert(0, str(PROJECT_ROOT))
+
+from .lib.model.unet1d import UNet1D
+from .lib.model.flow_matching import FlowMatching
+
+class PoseSampler:
+    def __init__(self, 
+                 device='cpu', 
+                 dataset='pp'):
+        """
+        初始化 PoseSampler 类。
+        在此处加载配置、统计量和模型，确保这些重型操作只执行一次。
+        """
+        self.device = torch.device(device)
+        self.checkpoint_path = "sample_utils/lib/ckpt/fm_20251117_172225/checkpoints/best.pt"
+        self.dataset = dataset
+        self.stats_dir = "sample_utils/lib/data_stats"
+        
+        # 1. 自动检测 run_dir 并加载配置
+        self.run_dir = self._detect_run_dir(self.checkpoint_path)
+        self.config = self._load_config(self.run_dir)
+        
+        # 2. 加载数据统计量 (Mean/Std)
+        self.pose_mean, self.pose_std = self._load_pose_stats()
+        
+        # 3. 加载模型
+        self.model = self._load_model()
+        print(f"Pose Sampler loaded successfully from {self.checkpoint_path}")
+
+    def _detect_run_dir(self, checkpoint_path):
+        """内部辅助方法：从 checkpoint 路径推断 run 目录"""
+        abs_path = os.path.abspath(checkpoint_path)
+        if 'checkpoints' in abs_path:
+            checkpoint_dir = os.path.dirname(abs_path)
+            if os.path.basename(checkpoint_dir) == 'checkpoints':
+                return os.path.dirname(checkpoint_dir)
+        return os.path.dirname(abs_path)
+
+    def _load_config(self, run_dir):
+        """内部辅助方法：加载 config.json"""
+        config_path = os.path.join(run_dir, 'config.json')
+        if not os.path.exists(config_path):
+            raise FileNotFoundError(f"Config file not found at {config_path}")
+            
+        with open(config_path, 'r') as f:
+            config = json.load(f)
+        return config
+
+    def _load_pose_stats(self):
+        """内部辅助方法：加载 Pose 统计量"""
+        if self.dataset == "pp":
+            filename = "pose_stats.pt"
+        else:
+            filename = "t_pose_stats.pt"
+            
+        file_path = os.path.join(self.stats_dir, filename)
+        if not os.path.exists(file_path):
+            raise FileNotFoundError(f"Stats file not found at {file_path}")
+
+        stats = torch.load(file_path, map_location=self.device)
+        return stats['mean'], stats['std']
+
+    def _load_model(self):
+        """内部辅助方法：初始化并加载模型权重"""
+        unet = UNet1D(
+            pose_dim=72,
+            base_channels=self.config['base_channels'],
+            channel_multipliers=self.config['channel_multipliers'],
+            time_emb_dim=self.config['time_emb_dim'],
+            mid_structure=self.config['mid_structure'],
+            mid_num_heads=self.config['mid_num_heads']
+        ).to(self.device)
+
+        model = FlowMatching(
+            model=unet,
+            sigma=0.0
+        ).to(self.device)
+
+        checkpoint = torch.load(self.checkpoint_path, map_location=self.device)
+
+        if 'model_state_dict' in checkpoint:
+            model.load_state_dict(checkpoint['model_state_dict'])
+        else:
+            model.load_state_dict(checkpoint)
+
+        model.eval()
+        return model
+
+    def _unnormalize_pose(self, pose_norm):
+        """内部辅助方法：反归一化"""
+        return pose_norm * self.pose_std + self.pose_mean
+
+    @torch.no_grad()
+    def sample(self, batch_size=1, num_steps=100, method="euler", verbose=False):
+        """
+        采样方法。
+        每次调用只需执行推理，无需重新加载模型。
+        """
+        samples_norm = self.model.sample(
+            sample_shape=(batch_size, 72),
+            device=self.device,
+            num_steps=num_steps,
+            method=method,
+            verbose=verbose
+        )
+        
+        samples_raw = self._unnormalize_pose(samples_norm)
+        return samples_raw
+
+# 使用示例
+if __name__ == '__main__':
+    # 1. 实例化 Sampler (只加载一次模型，耗时操作在这里)
+    print("Initializing sampler...")
+    sampler = PoseSampler(
+        device='cpu',
+        dataset='pp'
+    )
+    
+    # 2. 多次采样 (非常快)
+    print("Sampling batch 1...")
+    pose_batch_1 = sampler.sample(batch_size=1)
+    
+    print("Sampling batch 2 (with different batch size)...")
+    pose_batch_2 = sampler.sample(batch_size=4) # 甚至可以改变 batch size
+    
+    print(f"Batch 1 shape: {pose_batch_1.shape}")
+    print(f"Batch 2 shape: {pose_batch_2.shape}")
+    
+    import pdb; pdb.set_trace()
+    

src/sample_utils/trans_sampler.py

@@ -0,0 +1,135 @@
+import torch
+import numpy as np
+from scipy.stats import truncnorm
+
+
+def sample_transl4pp(batch_size, device):
+    """
+    为 SMPL 模型生成全局平移参数 (transl)，与数据集分布相似。
+
+    参数：
+        batch_size (int): 采样数量
+        device (torch.device): 计算设备 (e.g., 'cuda' or 'cpu')
+
+    返回：
+        transl (torch.Tensor): 形状为 (batch_size, 3) 的全局平移参数 [X, Y, Z]
+    """
+    # X 坐标：均匀分布在 [0.45, 0.85]
+    x_min, x_max = 0.45, 0.85
+    x = torch.rand(batch_size, 1) * (x_max - x_min) + x_min  # 均匀分布采样
+    
+    # Y 坐标：均匀分布在 [1.05, 1.45]
+    y_min, y_max = 1.05, 1.45
+    y = torch.rand(batch_size, 1) * (y_max - y_min) + y_min  # 均匀分布采样
+    
+    # 卧姿对应
+    # Z 坐标：截断正态分布，均值 0.08，标准差 0.03，范围 [-0.02, 0.24]
+    z_mean, z_std = 0.08, 0.03
+    z_min, z_max = -0.02, 0.24
+    # 计算截断正态分布的标准化边界
+    a, b = (z_min - z_mean) / z_std, (z_max - z_mean) / z_std
+    # 使用 scipy 的 truncnorm 生成截断正态分布采样
+    z = truncnorm.rvs(a, b, loc=z_mean, scale=z_std, size=(batch_size, 1))
+    z = torch.tensor(z, dtype=torch.float32)
+    
+    # # 站姿适应
+    # # Z 坐标：均匀分布在 [0.75, 0.85]
+    # z_min, z_max = 0.75, 0.85
+    # z = torch.rand(batch_size, 1) * (z_max - z_min) + z_min  # 均匀分布采样
+    
+    # 组合 X, Y, Z
+    transl = torch.cat([x, y, z], dim=1).to(device)
+    return transl
+
+
+def sample_transl4m(batch_size, device):
+    """
+    为 SMPL 模型生成全局平移参数 (transl)，与第二个数据集分布相似。
+    
+    参数：
+        batch_size (int): 采样数量
+        device (torch.device): 计算设备 (e.g., 'cuda' or 'cpu')
+    
+    返回：
+        transl (torch.Tensor): 形状为 (batch_size, 3) 的全局平移参数 [X, Y, Z]
+    """
+    # X 坐标：截断正态分布，均值 0.030337209，标准差 0.059348222，范围 [-0.22423534, 0.3106258]
+    x_mean, x_std = 0.030337209, 0.059348222
+    # x_min, x_max = -0.22423534, 0.3106258
+    x_min, x_max = -0.05, 0.1
+    # 计算截断正态分布的标准化边界
+    a_x, b_x = (x_min - x_mean) / x_std, (x_max - x_mean) / x_std
+    # 使用 scipy 的 truncnorm 生成截断正态分布采样
+    x = truncnorm.rvs(a_x, b_x, loc=x_mean, scale=x_std, size=(batch_size, 1))
+    x = torch.tensor(x, dtype=torch.float32)
+
+    # Y 坐标：截断正态分布，均值 0.5841795，标准差 0.2390917，范围 [-0.09659827, 1.2293766]
+    y_mean, y_std = 0.5841795, 0.2390917
+    # y_min, y_max = -0.09659827, 1.2293766
+    y_min, y_max = 0.0, 1.2
+    # 计算截断正态分布的标准化边界
+    a_y, b_y = (y_min - y_mean) / y_std, (y_max - y_mean) / y_std
+    # 使用 scipy 的 truncnorm 生成截断正态分布采样
+    y = truncnorm.rvs(a_y, b_y, loc=y_mean, scale=y_std, size=(batch_size, 1))
+    y = torch.tensor(y, dtype=torch.float32)
+
+    # # Z 坐标：均匀分布在 [0.75, 0.85]
+    # z_min, z_max = 0.75, 0.85
+    # z = torch.rand(batch_size, 1) * (z_max - z_min) + z_min  # 均匀分布采样
+    
+    # z全零
+    z = torch.zeros((batch_size, 1))
+    
+    # 组合 X, Y, Z
+    transl = torch.cat([x, y, z], dim=1).to(device)
+    return transl
+
+
+def sample_transl4t(batch_size, device):
+    """
+    为 SMPL 模型生成全局平移参数 (transl)，与第三个数据集分布相似。
+
+    参数：
+        batch_size (int): 采样数量
+        device (torch.device): 计算设备 (e.g., 'cuda' or 'cpu')
+
+    返回：
+        transl (torch.Tensor): 形状为 (batch_size, 3) 的全局平移参数 [X, Y, Z]
+    """
+    # X 坐标：截断正态分布，均值 0.35497144，标准差 0.08321648，范围 [0.10660601, 0.72766024]
+    x_mean, x_std = 0.35497144, 0.08321648
+    # x_min, x_max = 0.10660601, 0.72766024
+    x_min, x_max = 0.15, 0.55
+    # 计算截断正态分布的标准化边界
+    a_x, b_x = (x_min - x_mean) / x_std, (x_max - x_mean) / x_std
+    # 使用 scipy 的 truncnorm 生成截断正态分布采样
+    x = truncnorm.rvs(a_x, b_x, loc=x_mean, scale=x_std, size=(batch_size, 1))
+    x = torch.tensor(x, dtype=torch.float32)
+    
+    # Y 坐标：截断正态分布，均值 0.943629，标准差 0.0685662，范围 [0.7616181, 1.4328215]
+    y_mean, y_std = 0.943629, 0.0685662
+    # y_min, y_max = 0.7616181, 1.4328215
+    y_min, y_max = 0.8, 1.1
+    # 计算截断正态分布的标准化边界
+    a_y, b_y = (y_min - y_mean) / y_std, (y_max - y_mean) / y_std
+    # 使用 scipy 的 truncnorm 生成截断正态分布采样
+    y = truncnorm.rvs(a_y, b_y, loc=y_mean, scale=y_std, size=(batch_size, 1))
+    y = torch.tensor(y, dtype=torch.float32)
+    
+    # Z 坐标：截断正态分布，均值 -0.15257776，标准差 0.055761524，范围 [-0.44515115, 0.021567477]
+    z_mean, z_std = -0.15257776, 0.055761524
+    # z_min, z_max = -0.44515115, 0.021567477
+    z_min, z_max = -0.18, -0.1
+    # 计算截断���态分布的标准化边界
+    a_z, b_z = (z_min - z_mean) / z_std, (z_max - z_mean) / z_std
+    # 使用 scipy 的 truncnorm 生成截断正态分布采样
+    z = truncnorm.rvs(a_z, b_z, loc=z_mean, scale=z_std, size=(batch_size, 1))
+    z = torch.tensor(z, dtype=torch.float32)
+    
+    # 组合 X, Y, Z
+    transl = torch.cat([x, y, z], dim=1).to(device)
+    return transl
+
+
+
+