BiliSakura/BitDance-Tokenizer-diffusers

1from __future__ import annotations2 3import torch4from torch import nn5 6from diffusers.configuration_utils import ConfigMixin, register_to_config7from diffusers.models.modeling_utils import ModelMixin8from transformers.activations import ACT2FN9 10 11class BitDanceProjector(ModelMixin, ConfigMixin):12    @register_to_config13    def __init__(14        self,15        in_dim: int,16        out_dim: int,17        hidden_act: str = "gelu_pytorch_tanh",18    ) -> None:19        super().__init__()20        self.activation_fn = ACT2FN[hidden_act]21        self.fc1 = nn.Linear(in_dim, out_dim)22        self.fc2 = nn.Linear(out_dim, out_dim)23 24    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:25        hidden_states = hidden_states.to(self.fc1.weight.dtype)26        hidden_states = self.fc1(hidden_states)27        hidden_states = self.activation_fn(hidden_states)28        hidden_states = self.fc2(hidden_states)29        return hidden_states30