"""Transformers configuration for the ONNX Runtime TeraTTS release."""

from transformers import PretrainedConfig


class TeraTTSConfig(PretrainedConfig):
    model_type = "teratts_onnx"

    def __init__(
        self,
        sample_rate: int = 44_100,
        default_diffusion_model: str = "distilled",
        default_voice: str = "ru_f1",
        voices: list[str] | None = None,
        automatic_russian_stress: bool = True,
        cross_language_prompt_note: str = (
            "For an English reference voice speaking Russian, try duration_scale below 1.0."
        ),
        **kwargs,
    ) -> None:
        super().__init__(**kwargs)
        self.sample_rate = sample_rate
        self.default_diffusion_model = default_diffusion_model
        self.default_voice = default_voice
        self.voices = voices or []
        self.automatic_russian_stress = automatic_russian_stress
        self.cross_language_prompt_note = cross_language_prompt_note