| """VoXtream2-RU — Gradio-демо (аналог HF space herimor/voxtream2, но с русской моделью). |
| |
| Работает в двух режимах: |
| 1) Локально: python ru_finetune/ft4/space/app.py --model-dir ru_finetune/ft4/infer_model |
| 2) HF Space: переменная окружения MODEL_REPO=<user>/voxtream2-ru (repo с файлами |
| model.safetensors, config.json, phoneme_to_token.json, ru_tokens.json) — |
| файлы скачиваются через hf_hub_download. |
| |
| Отличия от EN-демо: RUAccent ставит ударения (решает омографы за́мок/замо́к), |
| espeak-ru фонемизация, спец-токены расширенного словаря. |
| """ |
|
|
| import argparse |
| import json |
| import os |
| import re |
| import sys |
| from pathlib import Path |
|
|
| HERE = Path(__file__).resolve().parent |
| |
| if str(HERE) not in sys.path: |
| sys.path.insert(0, str(HERE)) |
| from v10_fix_dict import FIX as INTERJ_FIX, redup_phones as _interj_redup |
|
|
| LOCAL_REPO_ID = "LOCAL_RU" |
| VOWEL = set("aeiouyɑɛɔʌəɵɨæøœɐɒʉʊɪ") |
| RU_VOWELS = set("аеёиоуыэюя") |
| PUNCT = (".", ",", "?", "!") |
| |
| |
| |
| PROCLITIC = {"в": "v", "с": "s", "к": "k", "ж": "ʒ", "б": "b"} |
| STRIP_WORD = ".,?!—-«»\"'()…:;–„“”+" |
| |
| NON_PHONE = ".,?!—-«»\"'()…:;–„“”" |
| |
| SIL_AFTER = ".,!?" |
|
|
| |
| |
| _PUNCT_MAP = {":": ",", ";": ",", "—": ",", "–": ",", "…": ".", "«": "", "»": "", |
| "„": "", "“": "", "”": "", '"': "", "(": "", ")": ""} |
|
|
|
|
| TARGET_LUFS = -23.0 |
|
|
|
|
| def normalize_lufs(wav, sr): |
| """v10: промпт -> -23 LUFS. Тихая бытовая запись для модели OOD; корпус |
| нормализован к тому же уровню (стандарт NeMo/NVIDIA voice cloning).""" |
| import numpy as np |
| import pyloudnorm |
|
|
| mono = wav.mean(dim=0).numpy().astype("float32") |
| if len(mono) < int(0.5 * sr): |
| return wav |
| try: |
| loud = pyloudnorm.Meter(sr).integrated_loudness(mono) |
| except Exception: |
| return wav |
| if not np.isfinite(loud) or loud < -70: |
| return wav |
| out = wav * float(10 ** ((TARGET_LUFS - loud) / 20)) |
| peak = float(out.abs().max()) |
| if peak > 0.99: |
| out = out * (0.99 / peak) |
| return out |
|
|
|
|
| def condition_prompt(wav, sr): |
| """v10.1: подготовка границы промпта. Замер: 5/6 тестовых промптов обрезаны |
| ПОСРЕДИ слова (жёсткий срез на 8 с), и провалы тембра сидят в первых ~1 с |
| генерации (окно 0.5 с: sim 0.11-0.29), а единственный промпт с тихим концом |
| (sanya_agin) — единственный без переключений. В обучении модель всегда |
| продолжает ПОСЛЕ паузы с комнатным тоном. Делаем так же: |
| 1) если конец промпта — речь, режем назад до ближайшей тишины (>= 80 мс |
| ниже -35 dB от пика, в последних 2.5 с); |
| 2) добавляем 0.35 с комнатного тона, синтезированного из тихого окна |
| самого промпта (build_groups_v3.synth_tone). VOXTREAM_PROMPT_COND=0 выкл.""" |
| import numpy as np |
| if os.environ.get("VOXTREAM_PROMPT_COND", "1") != "1": |
| return wav |
| mono = wav.mean(dim=0).numpy().astype("float32") |
| n = len(mono) |
| fr = max(int(0.02 * sr), 1) |
| peak = float(np.abs(mono).max()) + 1e-9 |
| if n > int(3.0 * sr): |
| |
| |
| nf_all = n // fr |
| e_all = 20 * np.log10(np.sqrt((mono[: nf_all * fr].reshape(nf_all, fr) ** 2).mean(1)) + 1e-9) - 20 * np.log10(peak) |
| thr = min(float(np.percentile(e_all, 20)), -25.0) |
| win = int(min(4.0 * sr, n - 2.0 * sr)) |
| seg_e = e_all[-(win // fr):] |
| run, cut = 0, None |
| for i, q in enumerate(seg_e < thr): |
| run = run + 1 if q else 0 |
| if run >= 3: |
| cut = i |
| if cut is None: |
| |
| |
| tail_e = e_all[-int(1.2 * sr) // fr:] |
| deep = np.where(tail_e < float(np.median(e_all)) - 6.0)[0] |
| if len(deep): |
| cut = len(seg_e) - len(tail_e) + int(deep[-1]) + 1 |
| if cut is not None: |
| end = (nf_all - len(seg_e) + cut - 1) * fr |
| if end > int(2.0 * sr): |
| wav = wav[:, :end] |
| mono = mono[:end] |
| try: |
| import build_groups_v3 as B |
| tmpl = B.room_tone(mono) |
| tone = B.synth_tone(tmpl, int(0.35 * sr), seed=7).astype("float32") |
| except Exception: |
| tone = (np.random.randn(int(0.35 * sr)) * 1e-4).astype("float32") |
| import torch |
| tone_t = torch.from_numpy(tone)[None].expand(wav.shape[0], -1) |
| return torch.cat([wav, tone_t.to(wav.dtype)], dim=1) |
|
|
|
|
| def normalize_punct(text: str) -> str: |
| for src, dst in _PUNCT_MAP.items(): |
| text = text.replace(src, dst) |
| text = re.sub(r"\s+([,.!?])", r"\1", text) |
| text = re.sub(r"([,.!?])\1+", r"\1", text) |
| return re.sub(r"\s{2,}", " ", text).strip() |
|
|
|
|
| def is_vowel(t): |
| return any(c in VOWEL for c in t) |
|
|
|
|
| def stressed_idx(aw): |
| vi, cnt, s, i = -1, 0, aw.lower(), 0 |
| while i < len(s): |
| if s[i] == "+": |
| if i + 1 < len(s) and s[i + 1] in RU_VOWELS: |
| vi = cnt |
| i += 1 |
| continue |
| if s[i] in RU_VOWELS: |
| cnt += 1 |
| i += 1 |
| return vi |
|
|
|
|
| class RUAccentPhonemizer: |
| """Интерфейс ESpeak.phonemize + RUAccent-ударения (перенос ˈ на нужную гласную).""" |
|
|
| def __init__(self): |
| from ruaccent import RUAccent |
| from voxtream.utils.text.phonemizer import ESpeak |
| self.acc = RUAccent() |
| self.acc.load(omograph_model_size="turbo3.1", use_dictionary=True) |
| self.esp = ESpeak("ru") |
| self._runorm = None |
|
|
| def _normalize_digits(self, text: str) -> str: |
| """v10: цифры/числа -> слова (RUNorm, падежи/род учитывает). Пользователь |
| пишет «в 2024 году» — работает без ручной нормализации. «Ё», потерянную |
| RUNorm'ом («четвертом»), ниже восстановит RUAccent.""" |
| if not re.search(r"\d", text): |
| return text |
| if self._runorm is None: |
| from runorm import RUNorm |
| self._runorm = RUNorm() |
| self._runorm.load(model_size="medium", |
| workdir=str(HERE / "runorm_cache")) |
| |
| |
| try: |
| parts = re.split(r"(?<=[.!?…])\s+", text) |
| return " ".join( |
| self._runorm.norm(p) if re.search(r"\d", p) else p for p in parts |
| ) |
| except Exception as e: |
| print(f"[runorm] fail: {e}; текст без нормализации") |
| return text |
|
|
| def _accent(self, text: str) -> str: |
| """Ударения: ручной '+' перед гласной (зам+ок) имеет приоритет, |
| RUAccent ставит только в словах без ручной пометки.""" |
| if "+" not in text: |
| return self.acc.process_all(text) |
| plain = re.sub(r"\+", "", text) |
| auto = self.acc.process_all(plain) |
| manual_w, auto_w = text.split(), auto.split() |
| if len(manual_w) != len(auto_w): |
| return auto |
| return " ".join( |
| mw if "+" in mw else aw for mw, aw in zip(manual_w, auto_w) |
| ) |
|
|
| def phonemize(self, text, separator="|", language="ru"): |
| text = self._normalize_digits(text) |
| text = normalize_punct(text) |
| accented = self._accent(text) |
| clean = re.sub(r"\+", "", accented) |
| seq = self.esp.phonemize(clean, separator=separator, language="ru") |
| esp_words, acc_words = seq.split(), accented.split() |
| if len(esp_words) != len(acc_words): |
| return seq |
| out = [] |
| for ew, aw in zip(esp_words, acc_words): |
| phones = [p for p in ew.split(separator) if p] |
| |
| |
| bare = aw.strip(STRIP_WORD).lower() |
| if bare in PROCLITIC: |
| keep = PROCLITIC[bare] |
| tail_p = phones[-1][-1] if phones and phones[-1][-1] in "".join(PUNCT) else "" |
| phones = [keep + tail_p] if tail_p else [keep] |
| |
| |
| |
| elif (ij := INTERJ_FIX.get(bare) or _interj_redup(bare)) is not None: |
| tail_p = phones[-1][-1] if phones and phones[-1][-1] in "".join(PUNCT) else "" |
| phones = ij.split() |
| if tail_p: |
| phones[-1] += tail_p |
| tail = "" |
| if phones and phones[-1] and phones[-1][-1] in "".join(PUNCT): |
| tail = phones[-1][-1] |
| phones[-1] = phones[-1][:-1] |
| if not phones[-1]: |
| phones.pop() |
| |
| |
| |
| phones = [p for p in (q.strip(NON_PHONE) for q in phones) if p] |
| ti = stressed_idx(aw) |
| if ti >= 0 and phones: |
| cl = [p.replace("ˈ", "").replace("ˌ", "") for p in phones] |
| vp = [j for j, p in enumerate(cl) if is_vowel(p)] |
| if ti < len(vp): |
| cl[vp[ti]] = "ˈ" + cl[vp[ti]] |
| phones = cl |
| w = separator.join(phones) |
| out.append(w + tail if tail else w) |
| |
| |
| |
| |
| |
| if tail and tail in SIL_AFTER: |
| out.append("sil") |
| return " ".join(out) |
|
|
|
|
| def resolve_model_files(): |
| """-> dict имя_файла -> локальный путь (из --model-dir или MODEL_REPO).""" |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--model-dir", default=os.environ.get("MODEL_DIR", "")) |
| args, _ = ap.parse_known_args() |
|
|
| names = ["model.safetensors", "config.json", "phoneme_to_token.json", "ru_tokens.json"] |
| if args.model_dir: |
| d = Path(args.model_dir).resolve() |
| return {n: str(d / n) for n in names} |
| repo = os.environ.get("MODEL_REPO") |
| assert repo, "укажите --model-dir или env MODEL_REPO" |
| from huggingface_hub import hf_hub_download |
| return {n: hf_hub_download(repo, n) for n in names} |
|
|
|
|
| def main(): |
| files = resolve_model_files() |
| ru_tokens = json.load(open(files["ru_tokens.json"])) |
|
|
| |
| import voxtream.utils.generator.setup as S |
| import voxtream.utils.generator.text as T |
| from huggingface_hub import hf_hub_download as _real_hf |
|
|
| def _hf(repo_id, filename, **kw): |
| if repo_id == LOCAL_REPO_ID: |
| return files[filename] |
| return _real_hf(repo_id, filename, **kw) |
|
|
| S.hf_hub_download = _hf |
|
|
| _orig_ttp = T.text_to_phone_tokens |
|
|
| |
| |
| |
| q_prefix = os.environ.get("QUESTION_PREFIX", "0") == "1" |
|
|
| def _add_q_prefix(text: str) -> str: |
| out = [] |
| for sent in re.split(r"(?<=[.!?])\s+", str(text).strip()): |
| words = sent.split() |
| if sent.rstrip().endswith("?") and len(words) > 1: |
| words[0] += "?" |
| out.append(" ".join(words)) |
| return " ".join(out) |
|
|
| def _ttp(*a, **kw): |
| kw["normalize"] = False |
| kw["language"] = "ru" |
| if q_prefix and a and isinstance(a[0], str): |
| a = (_add_q_prefix(a[0]),) + a[1:] |
| elif q_prefix and "text" in kw: |
| kw["text"] = _add_q_prefix(kw["text"]) |
| return _orig_ttp(*a, **kw) |
|
|
| T.text_to_phone_tokens = _ttp |
|
|
| |
| |
| |
| import voxtream.utils.generator.prompt as PR |
| _orig_ta = PR.torchaudio |
|
|
| class _LufsTorchaudio: |
| def __getattr__(self, name): |
| return getattr(_orig_ta, name) |
|
|
| @staticmethod |
| def load(path, *a, **kw): |
| wav, sr = _orig_ta.load(path, *a, **kw) |
| return condition_prompt(normalize_lufs(wav, sr), sr), sr |
|
|
| PR.torchaudio = _LufsTorchaudio() |
|
|
| from voxtream.generator import SpeechGenerator |
| _orig_init = SpeechGenerator.__init__ |
|
|
| def _patched_init(self, *a, **kw): |
| _orig_init(self, *a, **kw) |
| self.ctx.phonemizer = RUAccentPhonemizer() |
|
|
| SpeechGenerator.__init__ = _patched_init |
|
|
| |
| |
| |
| |
| |
| |
| _orig_gs = SpeechGenerator.generate_stream |
| _PAUSE = {".": 0.39, "?": 0.42, "!": 0.41} |
|
|
| def _split_sentences(text: str): |
| parts = [p.strip() for p in re.split(r"(?<=[.!?…])\s+", text.strip()) if p.strip()] |
| return parts if len(parts) >= 2 else [text] |
|
|
| def _gs_split(self, prompt_audio_path, text, speaking_rate=None, enhance_prompt=None, |
| apply_vad=None, return_progress=False, min_streaming_rtf=None): |
| if os.environ.get("VOXTREAM_SENT_SPLIT", "1") != "1" or not isinstance(text, str): |
| yield from _orig_gs(self, prompt_audio_path, text, speaking_rate, enhance_prompt, |
| apply_vad, return_progress, min_streaming_rtf) |
| return |
| sents = _split_sentences(text) |
| if len(sents) < 2: |
| yield from _orig_gs(self, prompt_audio_path, text, speaking_rate, enhance_prompt, |
| apply_vad, return_progress, min_streaming_rtf) |
| return |
| import numpy as _np |
| sr = int(self.config.mimi_sr) |
| pos_off, time_off, last_prog = 0, 0.0, None |
| for k, sent in enumerate(sents): |
| last_pos, last_t = 0, 0.0 |
| for item in _orig_gs(self, prompt_audio_path, sent, speaking_rate, enhance_prompt, |
| apply_vad, return_progress, min_streaming_rtf): |
| if return_progress: |
| frame, gt, prog = item |
| prog = dict(prog) |
| last_pos = max(last_pos, int(prog.get("phone_position", 0) or 0)) |
| last_t = max(last_t, float(prog.get("time_sec", 0.0) or 0.0)) |
| prog["phone_position"] = pos_off + int(prog.get("phone_position", 0) or 0) |
| prog["time_sec"] = time_off + float(prog.get("time_sec", 0.0) or 0.0) |
| last_prog = prog |
| yield frame, gt, prog |
| else: |
| yield item |
| pos_off += last_pos + 1 |
| time_off += last_t |
| if k < len(sents) - 1: |
| gap = _PAUSE.get(sent[-1], 0.39) |
| n = int(gap * sr) |
| noise = (_np.random.randn(n) * 1e-4).astype(_np.float32) |
| time_off += gap |
| if return_progress: |
| prog = dict(last_prog or {}) |
| prog["time_sec"] = time_off |
| yield noise, 0.0, prog |
| else: |
| yield noise, 0.0 |
|
|
| |
| |
| |
| |
| |
| from voxtream.utils.generator.prompt import extract_speaker_template as _est |
| import voxtream.utils.generator.prompt as _PR |
|
|
| def _spk_score(self, prompt_audio_path, audio_np): |
| import numpy as _np |
| import torch as _t |
| wav, sr = _PR.torchaudio.load(str(prompt_audio_path)) |
| pe = _est(wav.mean(0, keepdim=True), sr, self.ctx.spk_enc, self.config.spk_enc_sr, |
| self.ctx.device, self.ctx.dtype).float().reshape(-1) |
| gen = _t.from_numpy(audio_np.astype("float32"))[None] |
| osr = int(self.config.mimi_sr) |
| W, H = int(1.5 * osr), int(0.5 * osr) |
| sims = [] |
| for st in range(0, max(gen.shape[1] - W, 1), H): |
| e = _est(gen[:, st:st + W], osr, self.ctx.spk_enc, self.config.spk_enc_sr, |
| self.ctx.device, self.ctx.dtype).float().reshape(-1) |
| sims.append(float(pe @ e)) |
| if not sims: |
| return 0.0, 0.0 |
| return float(min(sims)), float(_np.median(sims)) |
|
|
| def _gs_bestof(self, prompt_audio_path, text, speaking_rate=None, enhance_prompt=None, |
| apply_vad=None, return_progress=False, min_streaming_rtf=None): |
| try: |
| N = int(os.environ.get("VOXTREAM_BEST_OF", "2")) |
| except ValueError: |
| N = 1 |
| if N <= 1 or not isinstance(text, str): |
| yield from _gs_split(self, prompt_audio_path, text, speaking_rate, enhance_prompt, |
| apply_vad, return_progress, min_streaming_rtf) |
| return |
| import numpy as _np |
| best, best_key = None, None |
| for k in range(N): |
| items = list(_gs_split(self, prompt_audio_path, text, speaking_rate, enhance_prompt, |
| apply_vad, return_progress, min_streaming_rtf)) |
| audio = _np.concatenate([it[0] for it in items]) if items else _np.zeros(1, "float32") |
| mn, med = _spk_score(self, prompt_audio_path, audio) |
| key = (mn >= med - 0.25, mn) |
| print(f"[best-of-{N}] кандидат {k + 1}: sim_min {mn:.2f} med {med:.2f}", flush=True) |
| if best_key is None or key > best_key: |
| best, best_key = items, key |
| if best_key[0]: |
| break |
| yield from best |
|
|
| SpeechGenerator.generate_stream = _gs_bestof |
|
|
| |
| root = HERE |
| gen_cfg = json.load(open(root / "configs/generator.json")) |
| gen_cfg.update( |
| model_repo=LOCAL_REPO_ID, |
| unk_token=ru_tokens["unk"], eop_token=ru_tokens["unk"], |
| bos_token=ru_tokens["bos"], eos_token=ru_tokens["eos"], |
| sil_token=ru_tokens["sil"], |
| enhance_prompt=False, apply_vad=False, |
| |
| |
| |
| temperature=0.8, topk=50, |
| |
| |
| |
| |
| cfg_gamma=1.5, |
| ) |
| |
| |
| |
| |
| |
| |
| |
| |
| os.environ.setdefault("VOXTREAM_SPS_GAIN", "25") |
| |
| |
| |
| |
| os.environ.setdefault("VOXTREAM_TEMPO_MODE", "src") |
| |
| |
| |
| |
| |
| os.environ.setdefault("VOXTREAM_DWELL_CAPS", "cons=5,vow=12") |
|
|
| ru_gen = HERE / "generator_ru.json" |
| |
| overrides = os.environ.get("GEN_OVERRIDES") |
| if overrides: |
| gen_cfg.update(json.loads(overrides)) |
| print(f"[app] generator overrides: {overrides}") |
| json.dump(gen_cfg, open(ru_gen, "w"), indent=2) |
|
|
| examples = HERE / "examples_ru.json" |
| if not examples.exists(): |
| json.dump({"examples": []}, open(examples, "w")) |
|
|
| sys.argv = [ |
| "voxtream-app", |
| "--config", str(ru_gen), |
| "--app-config", str(p if (p := HERE / "app_ru.json").exists() |
| else root / "configs/app.json"), |
| |
| "--spk-rate-config", str( |
| p if (p := root / "configs/speaking_rate_ru.json").exists() |
| else root / "configs/speaking_rate.json" |
| ), |
| "--examples-config", str(examples), |
| ] |
| from voxtream.app import main as app_main |
| app_main() |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|