| """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"] | |