from transformers import PretrainedConfig


class PixelDiTConfig(PretrainedConfig):
    model_type = "pixeldit"

    def __init__(
        self,
        in_channels=3,
        num_groups=24,
        hidden_size=1536,
        pixel_hidden_size=16,
        pixel_attn_hidden_size=1152,
        pixel_num_groups=16,
        patch_depth=14,
        pixel_depth=2,
        num_text_blocks=4,
        patch_size=16,
        txt_embed_dim=2304,
        txt_max_length=300,
        use_text_rope=True,
        text_rope_theta=10000.0,
        repa_encoder_index=-1,
        use_pixel_abs_pos=True,
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.in_channels            = in_channels
        self.num_groups             = num_groups
        self.hidden_size            = hidden_size
        self.pixel_hidden_size      = pixel_hidden_size
        self.pixel_attn_hidden_size = pixel_attn_hidden_size
        self.pixel_num_groups       = pixel_num_groups
        self.patch_depth            = patch_depth
        self.pixel_depth            = pixel_depth
        self.num_text_blocks        = num_text_blocks
        self.patch_size             = patch_size
        self.txt_embed_dim          = txt_embed_dim
        self.txt_max_length         = txt_max_length
        self.use_text_rope          = use_text_rope
        self.text_rope_theta        = text_rope_theta
        self.repa_encoder_index     = repa_encoder_index
        self.use_pixel_abs_pos      = use_pixel_abs_pos