from transformers import PretrainedConfig

class DiffusionConfig(PretrainedConfig):
    model_type = "gpt2"

    def __init__(
        self,
        vocab_size=425,
        d_model=512,
        n_heads=16,
        n_layers=6,
        dropout=0.1,
        max_len=2250,
        num_steps=1000,
        pad_token=0,
        **kwargs
    ):
        super().__init__(**kwargs)
        self.vocab_size = vocab_size
        self.d_model = d_model
        self.n_heads = n_heads
        self.n_layers = n_layers
        self.dropout = dropout
        self.max_len = max_len
        self.num_steps = num_steps
        self.pad_token = pad_token