ControlNet/marlin_vit_small_ytf

1from torch import nn, Tensor2from torch.nn import ModuleList, LayerNorm3 4from .modules import PatchEmbedding3d, Block5from .positional_embedding import SinCosPositionalEmbedding6 7 8class MarlinEncoder(nn.Module):9 10    def __init__(self, img_size=224, patch_size=16, n_frames=16, embed_dim=768, depth=12,11        num_heads=12, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop_rate=0., attn_drop_rate=0.,12        norm_layer="LayerNorm", init_values=0., tubelet_size=213    ):14        super().__init__()15 16        self.embed_dim = embed_dim17        self.patch_embedding = PatchEmbedding3d(18            input_size=(3, n_frames, img_size, img_size),19            patch_size=(tubelet_size, patch_size, patch_size),20            embedding=embed_dim21        )22        num_patches = (img_size // patch_size) * (img_size // patch_size) * (n_frames // tubelet_size)23 24        # sine-cosine positional embeddings25        self.pos_embedding = SinCosPositionalEmbedding((num_patches, embed_dim), dropout_rate=0.)26 27        if norm_layer == "LayerNorm":28            self.norm_layer = LayerNorm29            self.norm = self.norm_layer(embed_dim)30        else:31            raise NotImplementedError("Only LayerNorm is supported")32 33        self.blocks = ModuleList([34            Block(35                dim=embed_dim, num_heads=num_heads, mlp_ratio=mlp_ratio, qkv_bias=qkv_bias, qk_scale=qk_scale,36                drop=drop_rate, attn_drop=attn_drop_rate, norm_layer=self.norm_layer,37                init_values=init_values)38            for _ in range(depth)39        ])40 41        self.apply(self._init_weights)42 43    @staticmethod44    def _init_weights(m):45        if isinstance(m, nn.Linear):46            nn.init.xavier_uniform_(m.weight)47            if isinstance(m, nn.Linear) and m.bias is not None:48                nn.init.constant_(m.bias, 0)49        elif isinstance(m, nn.LayerNorm):50            nn.init.constant_(m.bias, 0)51            nn.init.constant_(m.weight, 1.0)52 53    def forward_features(self, x):54        for block in self.blocks:55            x = block(x)56        x = self.norm(x)57        return x58 59    def forward(self, x: Tensor, mask: Tensor) -> Tensor:60        # mask: (B, T, N) with boolean values, 0 -> masked, 1 -> visible61        assert len(x.shape) == 5, "x must be 5D"62        emb = self.patch_embedding(x)63        emb = self.pos_embedding(emb)64        b, _, c = emb.shape65        emb = emb[mask].view(b, -1, c)  # only visible patches are used66        emb = self.forward_features(emb)67        return emb68 69    def extract_features(self, x: Tensor, seq_mean_pool: bool) -> Tensor:70        x = self.patch_embedding(x)71        x = self.pos_embedding(x)72        for block in self.blocks:73            x = block(x)74 75        if seq_mean_pool:76            x = x.mean(dim=1)77        x = self.norm(x)78        return x79