Spaces:
Running on Zero
Running on Zero
| """ | |
| Skeleton context patch embedding for projecting VAE-encoded skeleton latents | |
| into DiT token space. Inherits from nn.Conv3d to preserve weight key compatibility | |
| with the original UranusAdapter.vace_patch_embedding checkpoint. | |
| Migrated from uranus.models.controller.UranusAdapter.vace_patch_embedding. | |
| """ | |
| from torch import nn | |
| class VacePatchEmbedding(nn.Conv3d): | |
| """ | |
| Conv3d that maps skeleton context into DiT token space. | |
| Input: [N_CAM, in_dim, F, H, W] — skeleton latents (2*z_dim + 64 channels) | |
| Output: [N_CAM, dim, F, H/2, W/2] — tokens matching DiT patch_size spatial stride | |
| Inherits from nn.Conv3d so state_dict key ``vace_patch_embedding.weight`` | |
| matches the original checkpoint without nesting. | |
| """ | |
| def __init__(self, in_dim=96, dim=1536): | |
| super().__init__(in_dim, dim, kernel_size=(1, 2, 2), stride=(1, 2, 2)) | |
| class LightweightSkeletonPatchEmbedding(nn.Conv3d): | |
| def __init__(self, in_dim=768, dim=1536): | |
| super().__init__(in_dim, dim, kernel_size=(1, 2, 2), stride=(1, 2, 2)) | |