from transformers import PretrainedConfig

class TopAIImageConfig(PretrainedConfig):
    model_type = "topai_image_gen"

    def __init__(
        self,
        input_dim=768,
        hidden_dim=512, # חזרנו ל-512 כדי להתאים למשקולות הקיימות
        image_channels=3,
        resolution=128,
        **kwargs
    ):
        super().__init__(**kwargs)
        self.input_dim = input_dim
        self.hidden_dim = hidden_dim
        self.image_channels = image_channels
        self.resolution = resolution