| """IndicF5-Hinglish -- self-contained loader.""" |
|
|
| import json |
| import os |
|
|
| import numpy as np |
| import torch |
| from huggingface_hub import hf_hub_download |
| from safetensors.torch import load_file |
| from transformers import PretrainedConfig, PreTrainedModel |
|
|
| from .cfm import CFM |
| from .dit import DiT |
| from .utils import get_tokenizer |
|
|
| SAMPLE_RATE = 24_000 |
| HOP_LENGTH = 256 |
| TARGET_RMS = 0.1 |
|
|
| PREFIXES = ("ema_model._orig_mod.", "ema_model.", "_orig_mod.", "model.", "module.") |
| IGNORE = {"initted", "step", |
| "mel_spec.mel_stft.mel_scale.fb", |
| "mel_spec.mel_stft.spectrogram.window"} |
|
|
|
|
| class IndicF5HinglishConfig(PretrainedConfig): |
| model_type = "indicf5_hinglish" |
|
|
| def __init__(self, dim=1024, depth=22, heads=16, ff_mult=2, text_dim=512, |
| conv_layers=4, mel_dim=100, **kw): |
| self.dim, self.depth, self.heads = dim, depth, heads |
| self.ff_mult, self.text_dim = ff_mult, text_dim |
| self.conv_layers, self.mel_dim = conv_layers, mel_dim |
| super().__init__(**kw) |
|
|
|
|
| def _strip_prefix(sd, target_keys): |
| variants = {"": sd} |
| for p in PREFIXES: |
| s = {k[len(p):]: v for k, v in sd.items() if k.startswith(p)} |
| if s: |
| variants[p] = s |
| prefix, best = max(variants.items(), key=lambda kv: len(set(kv[1]) & target_keys)) |
| return prefix, {k: v for k, v in best.items() if k not in IGNORE} |
|
|
|
|
| class IndicF5Hinglish(PreTrainedModel): |
| config_class = IndicF5HinglishConfig |
|
|
| def __init__(self, config, vocab_path=None, asset_dir=None): |
| super().__init__(config) |
| self.asset_dir = asset_dir |
| if vocab_path is None: |
| vocab_path = os.path.join(asset_dir or ".", "vocab.txt") |
|
|
| vocab_char_map, vocab_size = get_tokenizer(vocab_path, tokenizer="custom") |
| self.model = CFM( |
| transformer=DiT(dim=config.dim, depth=config.depth, heads=config.heads, |
| ff_mult=config.ff_mult, text_dim=config.text_dim, |
| conv_layers=config.conv_layers, |
| text_num_embeds=vocab_size, mel_dim=config.mel_dim), |
| mel_spec_kwargs=dict(n_fft=1024, hop_length=HOP_LENGTH, win_length=1024, |
| n_mel_channels=config.mel_dim, |
| target_sample_rate=SAMPLE_RATE, |
| mel_spec_type="vocos"), |
| odeint_kwargs=dict(method="euler"), |
| vocab_char_map=vocab_char_map, |
| ) |
| self._vocoder = None |
|
|
| @classmethod |
| def from_pretrained(cls, path, device=None, **kw): |
| device = device or ("cuda" if torch.cuda.is_available() else "cpu") |
|
|
| if os.path.isdir(path): |
| get = lambda f: os.path.join(path, f) |
| asset_dir = path |
| else: |
| get = lambda f: hf_hub_download(path, filename=f) |
| asset_dir = os.path.dirname(get("vocab.txt")) |
| vj = get("voices/voices.json") |
| for _v in json.load(open(vj)).values(): |
| get("voices/" + _v["file"]) |
|
|
| cfg_dict = json.load(open(get("config.json"))) |
| config = IndicF5HinglishConfig(**cfg_dict.get("model_config", {})) |
| model = cls(config, vocab_path=get("vocab.txt"), asset_dir=asset_dir) |
|
|
| sd = load_file(get("model.safetensors")) |
| target = set(model.model.state_dict()) |
| prefix, sd = _strip_prefix(sd, target) |
|
|
| missing, unexpected = model.model.load_state_dict(sd, strict=False) |
| real = [k for k in missing if k not in IGNORE] |
| loaded = len(target) - len(real) |
| if loaded < 0.9 * len(target): |
| raise RuntimeError(f"only {loaded}/{len(target)} tensors loaded " |
| f"(prefix={prefix!r}); missing e.g. {real[:5]}") |
| print(f"loaded {loaded}/{len(target)} tensors (prefix={prefix!r})") |
| return model.to(device).eval() |
|
|
| @property |
| def vocoder(self): |
| if self._vocoder is None: |
| from vocos import Vocos |
| self._vocoder = Vocos.from_pretrained( |
| "charactr/vocos-mel-24khz").to(self.device).eval() |
| return self._vocoder |
|
|
| def voices(self): |
| """Available bundled voices, keyed by name.""" |
| return json.load(open(os.path.join(self.asset_dir, "voices", "voices.json"))) |
|
|
| def voice(self, key): |
| """Returns (audio_path, transcript) for a bundled voice.""" |
| meta = self.voices() |
| if key not in meta: |
| raise KeyError(f"{key!r} not found; available: {list(meta)}") |
| v = meta[key] |
| return os.path.join(self.asset_dir, "voices", v["file"]), v["transcript"] |
|
|
| def default_voice(self): |
| meta = json.load(open(os.path.join(self.asset_dir, "voices", "voices.json"))) |
| v = meta[next(iter(meta))] |
| return os.path.join(self.asset_dir, "voices", v["file"]), v["transcript"] |
|
|
| def _load_ref(self, path): |
| import soundfile as sf |
| import librosa |
| audio, sr = sf.read(path, dtype="float32") |
| if audio.ndim > 1: |
| audio = audio.mean(axis=1) |
| if sr != SAMPLE_RATE: |
| audio = librosa.resample(audio, orig_sr=sr, target_sr=SAMPLE_RATE) |
| t = torch.from_numpy(audio).unsqueeze(0) |
| rms = torch.sqrt(torch.mean(t ** 2)) |
| if rms < TARGET_RMS: |
| t = t * TARGET_RMS / rms |
| return t.to(self.device), rms |
|
|
| @torch.inference_mode() |
| def generate(self, text, ref_audio=None, ref_text=None, nfe_step=32, |
| cfg_strength=2.0, sway_sampling_coef=-1.0, speed=1.0): |
| if ref_audio is None or ref_text is None: |
| ref_audio, ref_text = self.default_voice() |
|
|
| cond, rms = self._load_ref(ref_audio) |
| ref_text = ref_text.strip() |
| if not ref_text.endswith((" ", ".", "।", "!", "?")): |
| ref_text += ". " |
|
|
| ref_len = cond.shape[-1] // HOP_LENGTH |
| duration = ref_len + int( |
| ref_len / len(ref_text.encode()) * len(text.encode()) / speed) |
|
|
| generated, _ = self.model.sample( |
| cond=cond, text=[ref_text + text], duration=duration, |
| steps=nfe_step, cfg_strength=cfg_strength, |
| sway_sampling_coef=sway_sampling_coef) |
|
|
| mel = generated.to(torch.float32)[:, ref_len:, :].permute(0, 2, 1) |
| wave = self.vocoder.decode(mel).squeeze().cpu() |
| if rms < TARGET_RMS: |
| wave = wave * rms / TARGET_RMS |
| return wave.numpy().astype(np.float32), SAMPLE_RATE |
|
|
