| import torch.nn as nn |
| from einops import rearrange |
|
|
|
|
| class ControlEmbedder(nn.Module): |
| """Zero-init patch-embed for the control (pose-skeleton) latent. |
| |
| Mirrors the base PatchEmbed (spatial_patch_size=2, temporal_patch_size=1) but on the |
| raw 16-channel control latent (no padding-mask channel). Zero-initialized so it adds |
| nothing until trained (ControlNet zero-init). Output matches prepare_embedded_sequence's |
| x_B_T_H_W_D: (B, T, H/2, W/2, model_channels). |
| """ |
|
|
| def __init__(self, in_channels=16, spatial_patch_size=2, temporal_patch_size=1, model_channels=2048): |
| super().__init__() |
| self.spatial_patch_size = spatial_patch_size |
| self.temporal_patch_size = temporal_patch_size |
| in_features = in_channels * spatial_patch_size * spatial_patch_size * temporal_patch_size |
| self.proj = nn.Linear(in_features, model_channels) |
| nn.init.zeros_(self.proj.weight) |
| nn.init.zeros_(self.proj.bias) |
|
|
| def forward(self, control_B_C_T_H_W): |
| x = rearrange( |
| control_B_C_T_H_W, |
| "b c (t r) (h m) (w n) -> b t h w (c r m n)", |
| r=self.temporal_patch_size, m=self.spatial_patch_size, n=self.spatial_patch_size, |
| ) |
| return self.proj(x) |
|
|
