"""Configuration for QaDiT — a latent Diffusion Transformer for text-to-audio."""

from __future__ import annotations

from transformers import PretrainedConfig


class QaDiTConfig(PretrainedConfig):
    """Configuration of :class:`QaDiTModel`.

    Args:
        latent_channels / latent_time / latent_freq: geometry of the AudioLDM
            KL-VAE latent the DiT denoises, ``[C, T, F]``.
        patch_size: side of the square patch used to tokenize the latent.
        hidden_size / depth / num_heads / mlp_ratio: DiT backbone width, depth
            and attention/MLP shape.
        repa_layer: block index whose hidden state was aligned with AST
            features during training. Inference never uses it; it is kept so a
            converted model can be fine-tuned with the same REPA recipe.
        text_dim / text_max_length: hidden size and padded length of the frozen
            T5 encoder states consumed by cross-attention.
        latent_scale: unit-variance scaling applied to latents at precompute
            time. Sampling divides by it before VAE decoding, so it must match
            the ``latent_scale`` in the training data's ``meta.json``.
        load_auxiliaries: whether ``generate`` may pull T5/VAE/vocoder from the
            Hub automatically.
    """

    model_type = "qadit"
    attribute_map = {"num_hidden_layers": "depth", "num_attention_heads": "num_heads"}

    def __init__(
        self,
        latent_channels: int = 8,
        latent_time: int = 256,
        latent_freq: int = 16,
        patch_size: int = 2,
        hidden_size: int = 768,
        depth: int = 12,
        num_heads: int = 12,
        mlp_ratio: float = 4.0,
        repa_layer: int = 4,
        text_dim: int = 1024,
        text_max_length: int = 64,
        num_train_timesteps: int = 1000,
        schedule: str = "cosine",
        logit_normal_mean: float = 0.0,
        logit_normal_std: float = 1.0,
        num_inference_steps: int = 50,
        guidance_scale: float = 4.0,
        latent_scale: float = 1.0,
        sample_rate: int = 16_000,
        text_model: str = "google/flan-t5-large",
        vae_model: str = "cvssp/audioldm-s-full-v2",
        vae_subfolder: str = "vae",
        vocoder_model: str = "cvssp/audioldm-s-full-v2",
        vocoder_subfolder: str = "vocoder",
        load_auxiliaries: bool = True,
        **kwargs,
    ):
        self.latent_channels = latent_channels
        self.latent_time = latent_time
        self.latent_freq = latent_freq
        self.patch_size = patch_size
        self.hidden_size = hidden_size
        self.depth = depth
        self.num_heads = num_heads
        self.mlp_ratio = mlp_ratio
        self.repa_layer = repa_layer

        self.text_dim = text_dim
        self.text_max_length = text_max_length

        self.num_train_timesteps = num_train_timesteps
        self.schedule = schedule
        self.logit_normal_mean = logit_normal_mean
        self.logit_normal_std = logit_normal_std
        self.num_inference_steps = num_inference_steps
        self.guidance_scale = guidance_scale

        self.latent_scale = latent_scale
        self.sample_rate = sample_rate

        self.text_model = text_model
        self.vae_model = vae_model
        self.vae_subfolder = vae_subfolder
        self.vocoder_model = vocoder_model
        self.vocoder_subfolder = vocoder_subfolder
        self.load_auxiliaries = load_auxiliaries

        super().__init__(**kwargs)


__all__ = ["QaDiTConfig"]