| 36 | |
| 37 | |
| 38 | class UVit2DModel(ModelMixin, AttentionMixin, ConfigMixin, PeftAdapterMixin): |
| 39 | _supports_gradient_checkpointing = True |
| 40 | |
| 41 | @register_to_config |
| 42 | def __init__( |
| 43 | self, |
| 44 | # global config |
| 45 | hidden_size: int = 1024, |
| 46 | use_bias: bool = False, |
| 47 | hidden_dropout: float = 0.0, |
| 48 | # conditioning dimensions |
| 49 | cond_embed_dim: int = 768, |
| 50 | micro_cond_encode_dim: int = 256, |
| 51 | micro_cond_embed_dim: int = 1280, |
| 52 | encoder_hidden_size: int = 768, |
| 53 | # num tokens |
| 54 | vocab_size: int = 8256, # codebook_size + 1 (for the mask token) rounded |
| 55 | codebook_size: int = 8192, |
| 56 | # `UVit2DConvEmbed` |
| 57 | in_channels: int = 768, |
| 58 | block_out_channels: int = 768, |
| 59 | num_res_blocks: int = 3, |
| 60 | downsample: bool = False, |
| 61 | upsample: bool = False, |
| 62 | block_num_heads: int = 12, |
| 63 | # `TransformerLayer` |
| 64 | num_hidden_layers: int = 22, |
| 65 | num_attention_heads: int = 16, |
| 66 | # `Attention` |
| 67 | attention_dropout: float = 0.0, |
| 68 | # `FeedForward` |
| 69 | intermediate_size: int = 2816, |
| 70 | # `Norm` |
| 71 | layer_norm_eps: float = 1e-6, |
| 72 | ln_elementwise_affine: bool = True, |
| 73 | sample_size: int = 64, |
| 74 | ): |
| 75 | super().__init__() |
| 76 | |
| 77 | self.encoder_proj = nn.Linear(encoder_hidden_size, hidden_size, bias=use_bias) |
| 78 | self.encoder_proj_layer_norm = RMSNorm(hidden_size, layer_norm_eps, ln_elementwise_affine) |
| 79 | |
| 80 | self.embed = UVit2DConvEmbed( |
| 81 | in_channels, block_out_channels, vocab_size, ln_elementwise_affine, layer_norm_eps, use_bias |
| 82 | ) |
| 83 | |
| 84 | self.cond_embed = TimestepEmbedding( |
| 85 | micro_cond_embed_dim + cond_embed_dim, hidden_size, sample_proj_bias=use_bias |
| 86 | ) |
| 87 | |
| 88 | self.down_block = UVitBlock( |
| 89 | block_out_channels, |
| 90 | num_res_blocks, |
| 91 | hidden_size, |
| 92 | hidden_dropout, |
| 93 | ln_elementwise_affine, |
| 94 | layer_norm_eps, |
| 95 | use_bias, |
no outgoing calls
no test coverage detected
searching dependent graphs…