| |
| """ |
| Kokoro Arabic: Test Inference |
| ============================== |
| Synthesizes Palestinian Arabic test sentences from a StyleTTS2 checkpoint. |
| |
| Unlike test_inference.py (German), this bypasses KPipeline: Arabic has no |
| Kokoro lang_code, so we G2P with misaki espeak-ng 'ar' and apply the same |
| phoneme fixups as scripts/prepare_arabic.py, then call KModel directly. |
| |
| Usage: |
| CUDA_VISIBLE_DEVICES= python3 scripts/test_inference_arabic.py \ |
| --checkpoint StyleTTS2/logs/kokoro-arabic/epoch_1st_00000.pth \ |
| --voicepack voices/ar_kore_e1.pt \ |
| --output-dir test_output/epoch1_ar --device cpu |
| """ |
|
|
| import argparse |
| import sys |
| from pathlib import Path |
|
|
| _repo_root = Path(__file__).resolve().parents[1] |
| _kokoro_submodule = _repo_root / "kokoro" |
| if _kokoro_submodule.exists() and str(_kokoro_submodule) not in sys.path: |
| sys.path.insert(0, str(_kokoro_submodule)) |
|
|
| sys.path.insert(0, str(_repo_root / "scripts")) |
| from prepare_arabic import PHONEME_FIXUPS, normalize_arabic |
| from test_inference import convert_checkpoint |
|
|
| TEST_SENTENCES = [ |
| |
| "مرحبا، كيف حالك اليوم؟ إن شاء الله تكون مبسوط.", |
| "بدي أروح ع السوق أشتري خضرة وفواكه للبيت.", |
| |
| "الصبح الساعة سبعة طلعت من البيت على الشغل.", |
| "حكيت مع عمي عن العرس اللي صار الأسبوع الماضي.", |
| |
| "هاد حقه خمسين شيكل بس، مش غالي.", |
| |
| "شو رأيك نطلع نتمشى شوي بعد العصر؟", |
| |
| "قلت لقاسم القهوة بالسوق قريبة من البنك القديم.", |
| ] |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser(description="Test Arabic Kokoro inference") |
| group = parser.add_mutually_exclusive_group(required=True) |
| group.add_argument("--checkpoint", help="StyleTTS2 checkpoint (.pth), converted automatically") |
| group.add_argument("--model", help="Already-converted Kokoro-format weights (.pth)") |
| parser.add_argument("--voicepack", required=True) |
| parser.add_argument("--config", default="training/config.json") |
| parser.add_argument("--output-dir", default="test_output/arabic") |
| parser.add_argument("--device", default="cpu", choices=["cpu", "cuda"]) |
| parser.add_argument("--speed", type=float, default=1.0) |
| args = parser.parse_args() |
|
|
| import numpy as np |
| import soundfile as sf |
| import torch |
| from kokoro import KModel |
| from misaki import espeak |
|
|
| out = Path(args.output_dir) |
| out.mkdir(parents=True, exist_ok=True) |
|
|
| if args.checkpoint: |
| model_path = convert_checkpoint( |
| args.checkpoint, str(out / "kokoro_arabic_converted.pth") |
| ) |
| else: |
| model_path = args.model |
|
|
| print(f"Loading KModel on {args.device}...") |
| kmodel = KModel( |
| repo_id="hexgrad/Kokoro-82M", config=args.config, model=model_path |
| ) |
| kmodel = kmodel.to(args.device).eval() |
|
|
| voice = torch.load(args.voicepack, map_location="cpu", weights_only=True) |
| print(f"Voicepack: {args.voicepack} {tuple(voice.shape)}") |
|
|
| g2p = espeak.EspeakG2P(language="ar") |
|
|
| for i, text in enumerate(TEST_SENTENCES): |
| print(f"\n[{i + 1}/{len(TEST_SENTENCES)}] {text}") |
| ps, _ = g2p(normalize_arabic(text, g2p)) |
| for old, new in PHONEME_FIXUPS.items(): |
| ps = ps.replace(old, new) |
| ps = ps.strip() |
| print(f" phonemes: {ps}") |
| if len(ps) > 510: |
| ps = ps[:510] |
| ref_s = voice[len(ps) - 1] |
| with torch.no_grad(): |
| audio = kmodel(ps, ref_s, args.speed, return_output=False) |
| audio = audio.cpu().numpy().squeeze() |
| wav_path = out / f"ar_test_{i + 1:02d}.wav" |
| sf.write(str(wav_path), audio, 24000) |
| print(f" saved: {wav_path} ({len(audio) / 24000:.1f}s)") |
|
|
| print(f"\nDone! Output in {out}/") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
