68 lines
2.4 KiB
Python
68 lines
2.4 KiB
Python
"""
|
|
Condition Builder (with cmask mechanism).
|
|
|
|
Concatenates image token + skeleton token into the condition stream.
|
|
cmask: learnable mask token; both training and inference use HaMeR confidence as a continuous mask.
|
|
|
|
Inputs:
|
|
- image_tokens: (B, T, d_model) — from Image Refiner + Compression
|
|
- skeleton_tokens: (B, T, d_model) — from Skeleton Encoder
|
|
- hamer_confidence: (B, T) — HaMeR detection confidence
|
|
|
|
Output: condition (B, 2*T, d_model)
|
|
"""
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
|
|
|
|
class ConditionBuilder(nn.Module):
|
|
def __init__(self, d_model: int, mask_ratio: float = 0.2):
|
|
super().__init__()
|
|
self.d_model = d_model
|
|
self.mask_ratio = mask_ratio
|
|
|
|
# Learnable mask token — initialized with truncated normal
|
|
self.c_mask = nn.Parameter(torch.empty(1, 1, d_model))
|
|
nn.init.trunc_normal_(self.c_mask, std=0.02)
|
|
|
|
def forward(
|
|
self,
|
|
image_tokens: torch.Tensor,
|
|
skeleton_tokens: torch.Tensor,
|
|
hamer_confidence: torch.Tensor,
|
|
) -> tuple[torch.Tensor, torch.Tensor]:
|
|
"""
|
|
Args:
|
|
image_tokens: (B, T, d_model)
|
|
skeleton_tokens: (B, T, d_model)
|
|
hamer_confidence: (B, T) — used as m_t during inference
|
|
|
|
Returns:
|
|
condition: (B, 2*T, d_model)
|
|
m: (B, T) — actual mask values used (for logging)
|
|
"""
|
|
B, T, D = image_tokens.shape
|
|
|
|
# concat: (B, T, 2, d_model) -> (B, 2*T, d_model)
|
|
tokens = torch.stack([image_tokens, skeleton_tokens], dim=2) # (B, T, 2, D)
|
|
tokens = tokens.reshape(B, 2 * T, D)
|
|
|
|
# Use HaMeR confidence as a continuous mask (consistent between training and inference)
|
|
m = hamer_confidence.float() # (B, T)
|
|
|
|
# Optional: randomly dropout some frames during training as regularization
|
|
if self.training and self.mask_ratio > 0:
|
|
keep = (torch.rand(B, T, device=tokens.device) >= self.mask_ratio).float()
|
|
m = m * keep
|
|
|
|
# Expand m to match 2*T tokens (the two tokens per frame share the same m)
|
|
# m: (B, T) -> (B, 2*T)
|
|
m_expanded = m.unsqueeze(-1).expand(-1, -1, 2).reshape(B, 2 * T) # (B, 2*T)
|
|
m_expanded = m_expanded.unsqueeze(-1) # (B, 2*T, 1)
|
|
|
|
# cmask: c'_t = m_t * c_t + (1 - m_t) * c_mask
|
|
condition = m_expanded * tokens + (1 - m_expanded) * self.c_mask
|
|
|
|
return condition, m
|