Mike0021/zonos2
3
1from __future__ import annotations2 3import json4import re5from dataclasses import dataclass6from functools import lru_cache7from pathlib import Path8from typing import Optional9 10from zonos2.utils.logger import init_logger11from zonos2.utils.moe_topk import normalize_special_topk_layers12 13logger = init_logger(__name__)14 15_HF_REPO_ID_RE = re.compile(r"^[\w.-]+/[\w.-]+$")16 17 18@lru_cache()19def resolve_model_path(model_path: str) -> str:20 """Resolve a model path to a local directory.21 22 Local paths are returned as-is. A Hugging Face repo id (e.g.23 ``Zyphra/ZONOS2``) is downloaded to the local HF cache first.24 """25 if Path(model_path).expanduser().exists():26 return model_path27 if _HF_REPO_ID_RE.match(model_path):28 from huggingface_hub import snapshot_download29 30 logger.info("Downloading checkpoint from Hugging Face: %s", model_path)31 return snapshot_download(32 model_path,33 allow_patterns=["*.json", "*.pth", "*.pt", "*.yaml"],34 )35 return model_path36 37 38@dataclass39class Zonos2Config:40 """Config class for zonos2 models loaded from training checkpoints.41 42 This mimics the HuggingFace config interface for compatibility with the server.43 """44 45 model_type: str = "zonos2"46 dtype: str = "bfloat16"47 48 # Model architecture49 n_layers: int = 850 dim: int = 51251 head_dim: int = 12852 n_heads: Optional[int] = None53 n_kv_heads: Optional[int] = None54 ffn_dim_multiplier: float = 4.055 multiple_of: int = 25656 norm_eps: float = 1e-557 rope_theta: float = 10000.058 max_seqlen: int = 204859 60 # TTS-specific61 n_codebooks: int = 962 codebook_size: int = 102463 eoa_id: int = 102464 audio_pad_id: int = 102565 text_vocab: Optional[int] = None66 speaker_enabled: bool = False67 speaker_embedding_dim: int = 12868 # Optional LDA projection applied to speaker embeddings before the speaker69 # projection. The projection weights live inside the model checkpoint.70 speaker_lda_dim: Optional[int] = None71 # Clean/noisy speaker-background marker tokens (2 ids at the text-vocab tail).72 speaker_background_token_enabled: bool = False73 # Accurate-mode marker token (1 id after the background markers).74 accurate_mode_token_enabled: bool = False75 speaking_rate_num_buckets: int = 076 speaking_rate_buckets: Optional[list[str]] = None77 quality_num_buckets: int = 078 quality_features: Optional[list[str]] = None79 quality_buckets: Optional[dict[str, list[str]]] = None80 quality_dropout: Optional[dict[str, float]] = None81 82 # MoE config83 moe_router_topk: int = 184 special_topk_layers: dict[int, int] | None = None85 moe_n_experts: int = 186 moe_router_dim: int = 12887 moe_start_from_layer: int = 088 moe_end_from_layer: int = 089 moe_impl: str = "grouped"90 moe_balancing_strategy: str = "legacy"91 92 # Loss93 loss_softcap: float = 15.094 95 def __post_init__(self) -> None:96 self.special_topk_layers = normalize_special_topk_layers(97 self.special_topk_layers98 )99 if self.speaking_rate_buckets is not None:100 self.speaking_rate_buckets = [str(item) for item in self.speaking_rate_buckets]101 if self.quality_features is not None:102 self.quality_features = [str(item) for item in self.quality_features]103 if self.quality_buckets is not None:104 self.quality_buckets = {105 str(feature): [str(item) for item in (buckets or [])]106 for feature, buckets in self.quality_buckets.items()107 }108 if self.quality_features is None:109 self.quality_features = list(self.quality_buckets.keys())110 if int(self.quality_num_buckets or 0) <= 0:111 self.quality_num_buckets = sum(112 len(self.quality_buckets.get(feature, ()))113 for feature in (self.quality_features or [])114 )115 if self.quality_dropout is not None:116 self.quality_dropout = {117 str(feature): float(dropout)118 for feature, dropout in self.quality_dropout.items()119 }120 self.speaker_background_token_enabled = bool(self.speaker_background_token_enabled)121 self.accurate_mode_token_enabled = bool(self.accurate_mode_token_enabled)122 123 def to_dict(self):124 return self.__dict__.copy()125 126 127def _load_zonos2_config(model_path: str) -> Zonos2Config:128 """Load config from zonos2 training checkpoint format.129 130 Supports:131 - config.yaml in parent directory (training run format)132 - params.json in checkpoint directory (release checkpoint format)133 """134 from pathlib import Path135 136 def _cfg_get(cfg, key: str, default=None):137 if cfg is None:138 return default139 if isinstance(cfg, dict):140 return cfg.get(key, default)141 return getattr(cfg, key, default)142 143 def _apply_data_sidecar(model_params: dict, data_cfg) -> None:144 if data_cfg is None:145 return146 147 rate_buckets = _cfg_get(data_cfg, "speaking_rate_buckets", None) or []148 rate_buckets = [str(item) for item in rate_buckets]149 if rate_buckets:150 model_params["speaking_rate_buckets"] = rate_buckets151 152 rate_count = int(model_params.get("speaking_rate_num_buckets") or 0)153 rate_enabled = bool(_cfg_get(data_cfg, "speaking_rate_enabled", False))154 if rate_enabled:155 sidecar_rate_count = len(rate_buckets)156 if sidecar_rate_count == 0:157 sidecar_rate_count = int(_cfg_get(data_cfg, "speaking_rate_num_buckets", 0) or 0)158 if sidecar_rate_count > 0:159 rate_count = sidecar_rate_count160 if not int(model_params.get("speaking_rate_num_buckets") or 0):161 model_params["speaking_rate_num_buckets"] = rate_count162 163 # Quality conditioning lives in the data section of training configs.164 if bool(_cfg_get(data_cfg, "quality_enabled", False)):165 raw_features = _cfg_get(data_cfg, "quality_features", None)166 if hasattr(raw_features, "items"):167 quality_features = [168 str(feature) for feature, enabled in raw_features.items() if bool(enabled)169 ]170 else:171 quality_features = [str(item) for item in (raw_features or ())]172 raw_buckets = _cfg_get(data_cfg, "quality_buckets", None) or {}173 quality_buckets = {174 str(feature): [str(item) for item in (raw_buckets.get(feature, None) or ())]175 for feature in (quality_features or raw_buckets.keys())176 }177 if quality_buckets and "quality_buckets" not in model_params:178 model_params["quality_buckets"] = quality_buckets179 model_params["quality_features"] = quality_features or list(180 quality_buckets.keys()181 )182 raw_dropout = _cfg_get(data_cfg, "quality_dropout", None)183 if raw_dropout is not None and "quality_dropout" not in model_params:184 if hasattr(raw_dropout, "items"):185 model_params["quality_dropout"] = {186 str(feature): float(dropout) for feature, dropout in raw_dropout.items()187 }188 189 # Training configs name the marker-token flags after their data pipeline.190 background_enabled = _cfg_get(data_cfg, "speaker_embedding_origin_token_enabled", None)191 if background_enabled is not None:192 model_params.setdefault(193 "speaker_background_token_enabled", bool(background_enabled)194 )195 accurate_enabled = _cfg_get(196 data_cfg, "speaker_embedding_cartesia_clone_source_token_enabled", None197 )198 if accurate_enabled is not None:199 model_params.setdefault("accurate_mode_token_enabled", bool(accurate_enabled))200 201 if rate_count > 0 and model_params.get("text_vocab") is None:202 try:203 from zonos2.tts.prompt import conditioned_text_vocab_size204 205 quality_count = sum(206 len(buckets)207 for buckets in (model_params.get("quality_buckets") or {}).values()208 )209 background_count = (210 2 if model_params.get("speaker_background_token_enabled") else 0211 )212 accurate_count = (213 1214 if model_params.get("accurate_mode_token_enabled") and background_count215 else 0216 )217 model_params["text_vocab"] = conditioned_text_vocab_size(218 rate_count, quality_count, background_count, accurate_count219 )220 logger.debug(221 "Resolved text_vocab=%d from data text vocabulary and conditioning buckets",222 model_params["text_vocab"],223 )224 except Exception as exc:225 logger.debug("Could not resolve conditioned text_vocab from data config: %s", exc)226 227 def _validate_model_type(model_params: dict) -> None:228 model_type = model_params.get("model_type")229 if model_type is not None and str(model_type) != "zonos2":230 raise ValueError(231 f"Unsupported model_type={model_type!r}. This release only loads zonos2 checkpoints."232 )233 234 model_path = Path(model_path)235 236 # Try params.json in checkpoint dir first237 params_json = model_path / "params.json"238 if params_json.exists():239 with open(params_json, "r") as f:240 params = json.load(f)241 # params.json may have nested structure242 if "model" in params:243 model_params = params["model"]244 else:245 model_params = params246 247 # Also check config.yaml for tokenizer sidecar fields stored in data.248 for parent in [model_path, model_path.parent, model_path.parent.parent]:249 config_yaml = parent / "config.yaml"250 if config_yaml.exists():251 try:252 from omegaconf import OmegaConf253 254 cfg = OmegaConf.load(config_yaml)255 _apply_data_sidecar(model_params, getattr(cfg, "data", None))256 except ImportError:257 import yaml258 259 with open(config_yaml, "r") as f:260 cfg = yaml.safe_load(f)261 _apply_data_sidecar(262 model_params, cfg.get("data") if isinstance(cfg, dict) else None263 )264 break265 266 _validate_model_type(model_params)267 result = Zonos2Config(268 **{k: v for k, v in model_params.items() if hasattr(Zonos2Config, k)}269 )270 logger.debug("Loaded Zonos2Config from params.json")271 return result272 273 # Try config.yaml in parent directories274 for parent in [model_path, model_path.parent, model_path.parent.parent]:275 config_yaml = parent / "config.yaml"276 logger.debug("Checking for config at: %s", config_yaml)277 if config_yaml.exists():278 logger.debug("Found config.yaml at: %s", config_yaml)279 try:280 from omegaconf import OmegaConf281 282 cfg = OmegaConf.load(config_yaml)283 # Navigate to model config284 if hasattr(cfg, "model"):285 model_cfg = OmegaConf.to_container(cfg.model, resolve=True)286 else:287 model_cfg = OmegaConf.to_container(cfg, resolve=True)288 # Also check data section for tokenizer sidecar fields.289 _apply_data_sidecar(model_cfg, getattr(cfg, "data", None))290 _validate_model_type(model_cfg)291 result = Zonos2Config(292 **{k: v for k, v in model_cfg.items() if hasattr(Zonos2Config, k)}293 )294 logger.debug("Loaded Zonos2Config from config.yaml")295 return result296 except ImportError:297 # omegaconf not available, try as regular yaml298 import yaml299 300 with open(config_yaml, "r") as f:301 cfg = yaml.safe_load(f)302 model_cfg = cfg.get("model", cfg)303 # Also check data section for tokenizer sidecar fields.304 _apply_data_sidecar(model_cfg, cfg.get("data") if isinstance(cfg, dict) else None)305 _validate_model_type(model_cfg)306 result = Zonos2Config(307 **{k: v for k, v in model_cfg.items() if hasattr(Zonos2Config, k)}308 )309 logger.debug("Loaded Zonos2Config from config.yaml (yaml)")310 return result311 312 raise ValueError(313 f"Could not find config.yaml or params.json for zonos2 checkpoint at {model_path}"314 )315 316 317def _is_zonos2_checkpoint(model_path: str) -> bool:318 """Check if path is a zonos2 checkpoint with a loadable config."""319 model_path = Path(model_path)320 321 # Check for params.json322 if (model_path / "params.json").exists():323 return True324 325 # Check for config.yaml in parent dirs (training run format)326 for parent in [model_path, model_path.parent, model_path.parent.parent]:327 if (parent / "config.yaml").exists():328 return True329 330 return False331 332 333@lru_cache()334def _load_config(model_path: str) -> Zonos2Config:335 model_path = resolve_model_path(model_path)336 if not _is_zonos2_checkpoint(model_path):337 raise ValueError(338 f"Unsupported checkpoint at {model_path!r}. This release only loads "339 "Zonos2 TTS checkpoints with config.yaml or params.json."340 )341 return _load_zonos2_config(model_path)342 343 344def cached_load_checkpoint_config(model_path: str) -> Zonos2Config:345 config = _load_config(model_path)346 return Zonos2Config(**config.to_dict())347