| from transformers.configuration_utils import PretrainedConfig | |
| class TinyImageGenConfig(PretrainedConfig): | |
| model_type = "tinyimagegen" | |
| def __init__( | |
| self, | |
| image_size: int = 32, | |
| in_channels: int = 3, | |
| patch_size: int = 4, | |
| hidden_size: int = 32, | |
| num_hidden_layers: int = 6, | |
| num_attention_heads: int = 4, | |
| num_key_value_heads: int = 2, | |
| intermediate_size: int = 48, | |
| swiglu_interval: int = 3, | |
| num_lanes: int = 4, | |
| use_xsa: bool = False, | |
| use_per_head_gating: bool = False, | |
| rope_theta: float = 2500.0, | |
| rms_norm_eps: float = 1e-5, | |
| initializer_range: float = 0.02, | |
| **kwargs, | |
| ): | |
| self.image_size = image_size | |
| self.in_channels = in_channels | |
| self.patch_size = patch_size | |
| self.hidden_size = hidden_size | |
| self.num_hidden_layers = num_hidden_layers | |
| self.num_attention_heads = num_attention_heads | |
| self.num_key_value_heads = num_key_value_heads | |
| self.intermediate_size = intermediate_size | |
| self.swiglu_interval = swiglu_interval | |
| self.num_lanes = num_lanes | |
| self.use_xsa = use_xsa | |
| self.use_per_head_gating = use_per_head_gating | |
| self.rope_theta = rope_theta | |
| self.rms_norm_eps = rms_norm_eps | |
| self.initializer_range = initializer_range | |
| self.head_dim = hidden_size // num_attention_heads | |
| self.num_patches_side = image_size // patch_size | |
| self.num_patches = self.num_patches_side ** 2 | |
| self.patch_dim = in_channels * (patch_size ** 2) | |
| self.auto_map = { | |
| "AutoConfig": "configuration_tinyimagegen.TinyImageGenConfig", | |
| "AutoModel": "modeling_tinyimagegen.TinyImageGenModelForImageDiffusion", | |
| "AutoModelForImageDiffusion": "modeling_tinyimagegen.TinyImageGenModelForImageDiffusion", | |
| } | |
| super().__init__(**kwargs) | |