Gabriel de Marmiesse
c90fc8ced6
The french 6 layers variant wasn't included because it's not good enough yet. Use the `french_24l` variant for now.
136 lines
3.0 KiB
Python
136 lines
3.0 KiB
Python
"""Configuration models for loading YAML config files."""
|
|
|
|
from pathlib import Path
|
|
|
|
import yaml
|
|
from pydantic import BaseModel, ConfigDict
|
|
|
|
CONFIGS_DIR = Path(__file__).parent.parent / "config"
|
|
|
|
|
|
class StrictModel(BaseModel):
|
|
model_config = ConfigDict(extra="forbid")
|
|
|
|
|
|
# Flow configuration
|
|
class FlowConfig(StrictModel):
|
|
dim: int
|
|
depth: int
|
|
|
|
|
|
# Transformer configuration for FlowLM
|
|
class FlowLMTransformerConfig(StrictModel):
|
|
hidden_scale: int
|
|
max_period: int
|
|
d_model: int
|
|
num_heads: int
|
|
num_layers: int
|
|
|
|
|
|
class LookupTable(StrictModel):
|
|
dim: int
|
|
n_bins: int
|
|
tokenizer: str
|
|
tokenizer_path: str
|
|
|
|
|
|
# Root configuration
|
|
class FlowLMConfig(StrictModel):
|
|
"""Root configuration model for YAML config files."""
|
|
|
|
dtype: str
|
|
|
|
# Nested configurations
|
|
flow: FlowConfig
|
|
transformer: FlowLMTransformerConfig
|
|
|
|
# conditioning
|
|
lookup_table: LookupTable
|
|
weights_path: str | None = None
|
|
insert_bos_before_voice: bool = False
|
|
|
|
|
|
# SEANet configuration
|
|
class SEANetConfig(StrictModel):
|
|
dimension: int
|
|
channels: int
|
|
n_filters: int
|
|
n_residual_layers: int
|
|
ratios: list[int]
|
|
kernel_size: int
|
|
residual_kernel_size: int
|
|
last_kernel_size: int
|
|
dilation_base: int
|
|
pad_mode: str
|
|
compress: int
|
|
|
|
|
|
# Transformer configuration for Mimi
|
|
class MimiTransformerConfig(StrictModel):
|
|
d_model: int
|
|
input_dimension: int
|
|
output_dimensions: tuple[int, ...]
|
|
num_heads: int
|
|
num_layers: int
|
|
layer_scale: float
|
|
context: int
|
|
max_period: float = 10000.0
|
|
dim_feedforward: int
|
|
|
|
|
|
# Quantizer configuration
|
|
class QuantizerConfig(StrictModel):
|
|
dimension: int
|
|
output_dimension: int
|
|
|
|
|
|
# Root configuration
|
|
class MimiConfig(StrictModel):
|
|
"""Root configuration model for Mimi YAML config files."""
|
|
|
|
dtype: str
|
|
|
|
# Sample rate and channels
|
|
sample_rate: int
|
|
channels: int
|
|
frame_rate: float
|
|
|
|
# SEANet configurations
|
|
seanet: SEANetConfig
|
|
|
|
# Transformer
|
|
transformer: MimiTransformerConfig
|
|
|
|
# Quantizer
|
|
quantizer: QuantizerConfig
|
|
weights_path: str | None = None
|
|
inner_dim: int | None = None
|
|
outer_dim: int | None = None
|
|
|
|
|
|
class Config(StrictModel):
|
|
flow_lm: FlowLMConfig
|
|
mimi: MimiConfig
|
|
weights_path: str | None = None
|
|
weights_path_without_voice_cloning: str | None = None
|
|
pad_with_spaces_for_short_inputs: bool = False
|
|
remove_semicolons: bool = False
|
|
model_recommended_frames_after_eos: int | None = None
|
|
|
|
|
|
def load_config(yaml_path: str | Path) -> Config:
|
|
yaml_path = Path(yaml_path)
|
|
|
|
if not yaml_path.exists():
|
|
if yaml_path.is_relative_to(CONFIGS_DIR):
|
|
raise FileNotFoundError(
|
|
f"Config file not found: {yaml_path}. "
|
|
f"Did you make a typo? Available languages: {[p.stem for p in CONFIGS_DIR.glob('*.yaml')]}"
|
|
)
|
|
raise FileNotFoundError(f"Config file not found: {yaml_path}. Did you make a typo?")
|
|
|
|
with open(yaml_path, "r") as f:
|
|
config_dict = yaml.safe_load(f)
|
|
|
|
return Config(**config_dict)
|