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