Sofelia TTS training checkpoints (Eliaa final + multispeaker + stage1)

README.md

@@ -0,0 +1,36 @@
+---
+license: apache-2.0
+tags:
+- tts
+- arabic
+- kokoro
+- styletts2
+- checkpoint
+---
+
+# Sofelia TTS — Training Checkpoints (private)
+
+Raw StyleTTS2 Stage-2 checkpoints for the Palestinian Arabic Kokoro fine-tune.
+These are training-state checkpoints (`{'net': {...module.prefixed...}, 'optimizer', 'epoch', ...}`),
+**not** the inference model. For inference use the public model
+[hamdallah/Sofelia-TTS-82M](https://huggingface.co/hamdallah/Sofelia-TTS-82M).
+
+## Checkpoints
+
+| File | What | Val |
+|---|---|---|
+| `checkpoints/sirin_eliaa_final_epoch2nd.pth` | **Eliaa** single-voice final (powers the public model) | 0.463 (Sirin val set) |
+| `checkpoints/multispeaker_final_epoch2nd.pth` | multispeaker base (Kore + Sirin) | 0.446 (mixed val set) |
+| `checkpoints/stage1_first_stage.pth` | Stage-1 acoustic checkpoint | — |
+
+## Provenance
+
+Kokoro-82M → Stage 1 (multispeaker, Palestinian Arabic) → Stage 2 (multispeaker)
+→ Stage 2 continuation on the single human voice (Eliaa/Sirin), stopped at the
+validation plateau / overfit point. `training/` holds the configs, the dialect
+frontend (`prepare_arabic.py`, `ar_lexicon.json`), and the voicepack/inference
+scripts needed to reproduce or extend.
+
+To convert a checkpoint to the public Kokoro inference format, strip the
+`module.` prefixes from each `net` component and load via `kokoro.KModel`
+(see `training/test_inference_arabic.py`).

checkpoints/multispeaker_final_epoch2nd.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:118c2baf22bb2098484162211aa68bc25f0b87442ccb18d167dcaeaacab7acaf
+size 1819000897

checkpoints/sirin_eliaa_final_epoch2nd.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:eedd8fc2c8f8eb3515934efec67e0ccb45265bdf1baee01aa34be59392c32633
+size 1819000897

checkpoints/stage1_first_stage.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7ef937d261aac6c6dba6c137d445e4d60efbcec731ad48a17083645c5b0d7afe
+size 1738415866

training/ar_lexicon.json

@@ -0,0 +1,30 @@
+{
+  "_comment": "Palestinian Arabic pronunciation lexicon: surface word -> diacritized rewrite fed to espeak. Add entries whenever a word sounds wrong; applied by prepare_arabic.normalize_arabic in both training prep and inference.",
+  "حركتلو": "حَرَكتلو",
+  "حركتله": "حَرَكتله",
+  "قيمتها": "قِيمِتْها",
+  "قيمته": "قِيمْتُه",
+  "سبعمية": "سَبْعَمِيِّه",
+  "ستمية": "سِتْمِيِّه",
+  "خمسمية": "خَمَسْمِيِّه",
+  "اربعمية": "أَرْبَعْمِيِّه",
+  "تلتمية": "تَلَتْمِيِّه",
+  "ميتين": "مِيتين",
+  "مبارح": "مْبارِح",
+  "سبعة": "سَبْعَه",
+  "هلق": "هَلَق",
+  "هلأ": "هَلَأ",
+  "هيك": "هِيك",
+  "وين": "وِين",
+  "تنين": "تْنين",
+  "قلت": "قُلْت",
+  "معلش": "مَعليش",
+  "معليش": "مَعليش",
+  "بتذكر": "بِتذَكر",
+  "بعرفش": "بَعْرَفِش",
+  "يخرب": "يِخْرِب",
+  "بدك": "بِدك",
+  "قديش": "قَديش",
+  "بستناك": "بَستَناك",
+  "زلمة": "زَلَمِه"
+}

training/config.json

@@ -0,0 +1,150 @@
+{
+  "istftnet": {
+    "upsample_kernel_sizes": [20, 12],
+    "upsample_rates": [10, 6],
+    "gen_istft_hop_size": 5,
+    "gen_istft_n_fft": 20,
+    "resblock_dilation_sizes": [
+      [1, 3, 5],
+      [1, 3, 5],
+      [1, 3, 5]
+    ],
+    "resblock_kernel_sizes": [3, 7, 11],
+    "upsample_initial_channel": 512
+  },
+  "dim_in": 64,
+  "dropout": 0.2,
+  "hidden_dim": 512,
+  "max_conv_dim": 512,
+  "max_dur": 50,
+  "multispeaker": true,
+  "n_layer": 3,
+  "n_mels": 80,
+  "n_token": 178,
+  "style_dim": 128,
+  "text_encoder_kernel_size": 5,
+  "plbert": {
+    "hidden_size": 768,
+    "num_attention_heads": 12,
+    "intermediate_size": 2048,
+    "max_position_embeddings": 512,
+    "num_hidden_layers": 12,
+    "dropout": 0.1
+  },
+  "vocab": {
+    ";": 1,
+    ":": 2,
+    ",": 3,
+    ".": 4,
+    "!": 5,
+    "?": 6,
+    "—": 9,
+    "…": 10,
+    "\"": 11,
+    "(": 12,
+    ")": 13,
+    "“": 14,
+    "”": 15,
+    " ": 16,
+    "\u0303": 17,
+    "ʣ": 18,
+    "ʥ": 19,
+    "ʦ": 20,
+    "ʨ": 21,
+    "ᵝ": 22,
+    "\uAB67": 23,
+    "A": 24,
+    "I": 25,
+    "O": 31,
+    "Q": 33,
+    "S": 35,
+    "T": 36,
+    "W": 39,
+    "Y": 41,
+    "ᵊ": 42,
+    "a": 43,
+    "b": 44,
+    "c": 45,
+    "d": 46,
+    "e": 47,
+    "f": 48,
+    "h": 50,
+    "i": 51,
+    "j": 52,
+    "k": 53,
+    "l": 54,
+    "m": 55,
+    "n": 56,
+    "o": 57,
+    "p": 58,
+    "q": 59,
+    "r": 60,
+    "s": 61,
+    "t": 62,
+    "u": 63,
+    "v": 64,
+    "w": 65,
+    "x": 66,
+    "y": 67,
+    "z": 68,
+    "ɑ": 69,
+    "ɐ": 70,
+    "ɒ": 71,
+    "æ": 72,
+    "β": 75,
+    "ɔ": 76,
+    "ɕ": 77,
+    "ç": 78,
+    "ɖ": 80,
+    "ð": 81,
+    "ʤ": 82,
+    "ə": 83,
+    "ɚ": 85,
+    "ɛ": 86,
+    "ɜ": 87,
+    "ɟ": 90,
+    "ɡ": 92,
+    "ɥ": 99,
+    "ɨ": 101,
+    "ɪ": 102,
+    "ʝ": 103,
+    "ɯ": 110,
+    "ɰ": 111,
+    "ŋ": 112,
+    "ɳ": 113,
+    "ɲ": 114,
+    "ɴ": 115,
+    "ø": 116,
+    "ɸ": 118,
+    "θ": 119,
+    "œ": 120,
+    "ɹ": 123,
+    "ɾ": 125,
+    "ɻ": 126,
+    "ʁ": 128,
+    "ɽ": 129,
+    "ʂ": 130,
+    "ʃ": 131,
+    "ʈ": 132,
+    "ʧ": 133,
+    "ʊ": 135,
+    "ʋ": 136,
+    "ʌ": 138,
+    "ɣ": 139,
+    "ɤ": 140,
+    "χ": 142,
+    "ʎ": 143,
+    "ʒ": 147,
+    "ʔ": 148,
+    "ˈ": 156,
+    "ˌ": 157,
+    "ː": 158,
+    "ʰ": 162,
+    "ʲ": 164,
+    "↓": 169,
+    "→": 171,
+    "↗": 172,
+    "↘": 173,
+    "ᵻ": 177
+  }
+}

training/config_arabic_ft.yml

@@ -0,0 +1,145 @@
+batch_size: 6  # stage 2 OOMs at 10 on A100-40GB (dual decode paths + WavLM); joint_epoch adds more
+epochs: 10
+epochs_1st: 10
+epochs_2nd: 10
+save_freq: 1
+# Stage 2 resume: weights from the epoch-1 stage-2 checkpoint (taa-marbuta phoneme
+# fix applied to lists on 2026-06-11; epoch counter restarts via load_only_params).
+# For a fresh stage-2 start, revert to second_stage_load_pretrained: false +
+# pretrained_model: "../training/kokoro_base.pth" (stage 1 reads pretrained_model too).
+pretrained_model: "logs/kokoro-arabic/epoch_2nd_00002_stripped.pth"  # epoch-3 ckpt (val 0.466); module. prefixes stripped — raw epoch_2nd_*.pth keys don't match load_checkpoint's bare modules
+first_stage_path: "first_stage.pth"
+load_only_params: true
+second_stage_load_pretrained: true
+# ─────────────────────────────────────────────────────────────────────────────
+# Kokoro Palestinian Arabic Fine-Tuning Config
+# ─────────────────────────────────────────────────────────────────────────────
+# Based on configs/config_german_ft.yml, adapted for:
+# - Two Palestinian Arabic voices:
+#     speaker 0 "Kore"  (~130h Gemini TTS story segments, /home/ubuntu/data)
+#     speaker 1 "Sirin" (Sofelia-Palestinian-TTS HF dataset)
+# - Train lists use ABSOLUTE wav paths (root_path only feeds TB voicepack previews)
+# - G2P: misaki espeak-ng 'ar' with fixups (see scripts/prepare_arabic.py)
+# ─────────────────────────────────────────────────────────────────────────────
+
+# ── Logging ──────────────────────────────────────────────────────────────────
+log_dir: "logs/kokoro-arabic"
+
+# ── Data ─────────────────────────────────────────────────────────────────────
+data_params:
+  train_data: "../training/train_list_ar.txt"
+  val_data: "../training/val_list_ar.txt"
+  # Train/val lists use ABSOLUTE wav paths (osp.join ignores root_path for them).
+  # root_path is only rglob'd by kokoro_tb_utils.extract_voicepack for TB previews —
+  # it must be a real audio dir, never "/" (rglob over the whole fs crashes on /proc).
+  root_path: "/home/ubuntu/data/audio"
+  OOD_data: "../training/OOD_texts_ar.txt"
+  min_length: 50                    # min phoneme sequence length (chars)
+  num_workers: 24                   # CPU workers for data loading
+
+# ── Preprocessing ────────────────────────────────────────────────────────────
+preprocess_params:
+  sr: 24000
+  spect_params:
+    n_fft: 2048
+    win_length: 1200
+    hop_length: 300
+    n_mels: 80
+    fmin: 0
+    fmax: 8000
+
+# ── Model Architecture ──────────────────────────────────────────────────────
+# Must match Kokoro-82M exactly to load pre-trained weights
+model_params:
+  dim_in: 64
+  n_token: 178                      # phoneme vocabulary size (IPA + punctuation)
+  hidden_dim: 512                   # model hidden dimension
+  style_dim: 128                    # style/speaker embedding dimension (x2 = 256 total in voicepack)
+  max_dur: 50                       # max duration per phoneme (in frames)
+  multispeaker: true                # two Arabic voices (Kore=0, Sirin=1)
+  n_mels: 80                        # mel spectrogram bands
+  dropout: 0.2
+  n_layer: 3                        # number of layers in text encoder / predictor
+  text_encoder_kernel_size: 5
+
+  decoder:
+    type: istftnet                  # Kokoro uses ISTFTNet, NOT hifigan
+    upsample_rates: [10, 6]
+    upsample_kernel_sizes: [20, 12]
+    upsample_initial_channel: 512
+    resblock_kernel_sizes: [3, 7, 11]
+    resblock_dilation_sizes: [[1, 3, 5], [1, 3, 5], [1, 3, 5]]
+    gen_istft_n_fft: 20
+    gen_istft_hop_size: 5
+
+  # PLBERT config (phoneme-level BERT for text understanding)
+  diffusion:
+    embedding_mask_proba: 0.1
+    transformer:
+      num_layers: 3
+      num_heads: 8
+      head_features: 64
+      multiplier: 2
+    dist:
+      sigma_data: 0.2
+      estimate_sigma_data: true
+      mean: -3.0
+      std: 1.0
+
+  plbert:
+    hidden_size: 768
+    num_attention_heads: 12
+    intermediate_size: 2048
+    max_position_embeddings: 512
+    num_hidden_layers: 12
+    dropout: 0.1
+
+  # SLM discriminator (WavLM) — used in stage 2 only
+  slm:
+    model: "microsoft/wavlm-base-plus"
+    sr: 16000
+    hidden: 768
+    nlayers: 13
+    initial_channel: 64
+
+# ── Loss Weights ─────────────────────────────────────────────────────────────
+loss_params:
+  lambda_gen: 1.0
+  lambda_mel: 5.0                   # mel reconstruction loss
+  lambda_dur: 1.0                   # duration prediction loss
+  lambda_ce: 20.0                   # duration cross-entropy loss
+  lambda_F0: 1.0                    # F0 (pitch) prediction loss
+  lambda_norm: 1.0                  # energy/norm prediction loss
+  lambda_s2s: 1.0                   # sequence-to-sequence loss
+  lambda_mono: 1.0                  # monotonic alignment loss
+  lambda_slm: 1.0                   # SLM adversarial (0 = disabled for stage 1)
+  lambda_diff: 0.0                  # diffusion loss (0 = disabled, no diffusion)
+  lambda_sty: 0.0                   # style reconstruction (0 = disabled)
+  TMA_epoch: 0                      # start alignment from epoch 0 (fine-tuning, not training from scratch)
+  diff_epoch: 999                   # effectively disable diffusion training
+  joint_epoch: 1                    # resumed run: 1 absorption epoch for frontend-v2 labels, then SLM joint training
+
+# ── Optimizer ────────────────────────────────────────────────────────────────
+optimizer_params:
+  lr: 0.0001                        # main learning rate
+  bert_lr: 0.00001                  # lower LR for PLBERT (mostly frozen knowledge)
+  ft_lr: 0.0001                     # fine-tuning LR for acoustic modules
+
+# ── Pitch Extractor ──────────────────────────────────────────────────────────
+F0_path: "Utils/JDC/bst.t7"
+
+# ── ASR Alignment Model ─────────────────────────────────────────────────────
+ASR_config: "Utils/ASR/config.yml"
+ASR_path: "Utils/ASR/epoch_00080.pth"
+
+# ── PLBERT ───────────────────────────────────────────────────────────────────
+PLBERT_dir: "Utils/PLBERT/"
+
+slmadv_params:
+  min_len: 100
+  max_len: 500
+  batch_percentage: 0.5
+  iter: 10
+  thresh: 5
+  scale: 0.01
+  sig: 1.5

training/config_sirin_ft.yml

@@ -0,0 +1,145 @@
+batch_size: 6  # stage 2 OOMs at 10 on A100-40GB (dual decode paths + WavLM); joint_epoch adds more
+epochs: 10
+epochs_1st: 10
+epochs_2nd: 10
+save_freq: 1
+# Stage 2 resume: weights from the epoch-1 stage-2 checkpoint (taa-marbuta phoneme
+# fix applied to lists on 2026-06-11; epoch counter restarts via load_only_params).
+# For a fresh stage-2 start, revert to second_stage_load_pretrained: false +
+# pretrained_model: "../training/kokoro_base.pth" (stage 1 reads pretrained_model too).
+pretrained_model: "logs/kokoro-arabic/epoch_2nd_00003_stripped.pth"  # epoch-3 ckpt (val 0.466); module. prefixes stripped — raw epoch_2nd_*.pth keys don't match load_checkpoint's bare modules
+first_stage_path: "first_stage.pth"
+load_only_params: true
+second_stage_load_pretrained: true
+# ─────────────────────────────────────────────────────────────────────────────
+# Kokoro Palestinian Arabic Fine-Tuning Config
+# ─────────────────────────────────────────────────────────────────────────────
+# Based on configs/config_german_ft.yml, adapted for:
+# - Two Palestinian Arabic voices:
+#     speaker 0 "Kore"  (~130h Gemini TTS story segments, /home/ubuntu/data)
+#     speaker 1 "Sirin" (Sofelia-Palestinian-TTS HF dataset)
+# - Train lists use ABSOLUTE wav paths (root_path only feeds TB voicepack previews)
+# - G2P: misaki espeak-ng 'ar' with fixups (see scripts/prepare_arabic.py)
+# ─────────────────────────────────────────────────────────────────────────────
+
+# ── Logging ──────────────────────────────────────────────────────────────────
+log_dir: "logs/kokoro-sirin"
+
+# ── Data ─────────────────────────────────────────────────────────────────────
+data_params:
+  train_data: "../training/train_list_sirin.txt"
+  val_data: "../training/val_list_sirin.txt"
+  # Train/val lists use ABSOLUTE wav paths (osp.join ignores root_path for them).
+  # root_path is only rglob'd by kokoro_tb_utils.extract_voicepack for TB previews —
+  # it must be a real audio dir, never "/" (rglob over the whole fs crashes on /proc).
+  root_path: "/home/ubuntu/data/audio"
+  OOD_data: "../training/OOD_texts_sirin.txt"
+  min_length: 50                    # min phoneme sequence length (chars)
+  num_workers: 24                   # CPU workers for data loading
+
+# ── Preprocessing ────────────────────────────────────────────────────────────
+preprocess_params:
+  sr: 24000
+  spect_params:
+    n_fft: 2048
+    win_length: 1200
+    hop_length: 300
+    n_mels: 80
+    fmin: 0
+    fmax: 8000
+
+# ── Model Architecture ──────────────────────────────────────────────────────
+# Must match Kokoro-82M exactly to load pre-trained weights
+model_params:
+  dim_in: 64
+  n_token: 178                      # phoneme vocabulary size (IPA + punctuation)
+  hidden_dim: 512                   # model hidden dimension
+  style_dim: 128                    # style/speaker embedding dimension (x2 = 256 total in voicepack)
+  max_dur: 50                       # max duration per phoneme (in frames)
+  multispeaker: true                # two Arabic voices (Kore=0, Sirin=1)
+  n_mels: 80                        # mel spectrogram bands
+  dropout: 0.2
+  n_layer: 3                        # number of layers in text encoder / predictor
+  text_encoder_kernel_size: 5
+
+  decoder:
+    type: istftnet                  # Kokoro uses ISTFTNet, NOT hifigan
+    upsample_rates: [10, 6]
+    upsample_kernel_sizes: [20, 12]
+    upsample_initial_channel: 512
+    resblock_kernel_sizes: [3, 7, 11]
+    resblock_dilation_sizes: [[1, 3, 5], [1, 3, 5], [1, 3, 5]]
+    gen_istft_n_fft: 20
+    gen_istft_hop_size: 5
+
+  # PLBERT config (phoneme-level BERT for text understanding)
+  diffusion:
+    embedding_mask_proba: 0.1
+    transformer:
+      num_layers: 3
+      num_heads: 8
+      head_features: 64
+      multiplier: 2
+    dist:
+      sigma_data: 0.2
+      estimate_sigma_data: true
+      mean: -3.0
+      std: 1.0
+
+  plbert:
+    hidden_size: 768
+    num_attention_heads: 12
+    intermediate_size: 2048
+    max_position_embeddings: 512
+    num_hidden_layers: 12
+    dropout: 0.1
+
+  # SLM discriminator (WavLM) — used in stage 2 only
+  slm:
+    model: "microsoft/wavlm-base-plus"
+    sr: 16000
+    hidden: 768
+    nlayers: 13
+    initial_channel: 64
+
+# ── Loss Weights ─────────────────────────────────────────────────────────────
+loss_params:
+  lambda_gen: 1.0
+  lambda_mel: 5.0                   # mel reconstruction loss
+  lambda_dur: 1.0                   # duration prediction loss
+  lambda_ce: 20.0                   # duration cross-entropy loss
+  lambda_F0: 1.0                    # F0 (pitch) prediction loss
+  lambda_norm: 1.0                  # energy/norm prediction loss
+  lambda_s2s: 1.0                   # sequence-to-sequence loss
+  lambda_mono: 1.0                  # monotonic alignment loss
+  lambda_slm: 1.0                   # SLM adversarial (0 = disabled for stage 1)
+  lambda_diff: 0.0                  # diffusion loss (0 = disabled, no diffusion)
+  lambda_sty: 0.0                   # style reconstruction (0 = disabled)
+  TMA_epoch: 0                      # start alignment from epoch 0 (fine-tuning, not training from scratch)
+  diff_epoch: 999                   # effectively disable diffusion training
+  joint_epoch: 1                    # resumed run: 1 absorption epoch for frontend-v2 labels, then SLM joint training
+
+# ── Optimizer ────────────────────────────────────────────────────────────────
+optimizer_params:
+  lr: 0.0001                        # main learning rate
+  bert_lr: 0.00001                  # lower LR for PLBERT (mostly frozen knowledge)
+  ft_lr: 0.0001                     # fine-tuning LR for acoustic modules
+
+# ── Pitch Extractor ──────────────────────────────────────────────────────────
+F0_path: "Utils/JDC/bst.t7"
+
+# ── ASR Alignment Model ─────────────────────────────────────────────────────
+ASR_config: "Utils/ASR/config.yml"
+ASR_path: "Utils/ASR/epoch_00080.pth"
+
+# ── PLBERT ───────────────────────────────────────────────────────────────────
+PLBERT_dir: "Utils/PLBERT/"
+
+slmadv_params:
+  min_len: 100
+  max_len: 500
+  batch_percentage: 0.5
+  iter: 10
+  thresh: 5
+  scale: 0.01
+  sig: 1.5

training/extract_voicepack.py

@@ -0,0 +1,467 @@
+#!/usr/bin/env python3
+"""
+Kokoro German: Extract Voicepack
+=================================
+Extracts a voicepack (.pt) from a fine-tuned StyleTTS2 checkpoint by running
+both style encoders (acoustic + prosodic) on representative utterances and
+averaging the resulting style vectors.
+
+Usage:
+    # Single checkpoint (Stage 1 only — uses style_encoder for both halves)
+    python scripts/extract_voicepack.py \
+        --model StyleTTS2/logs/kokoro_german/epoch_1st_00002.pth \
+        --audio-dir dataset/audio/dm_daniel \
+        --output voices/dm_daniel.pt
+
+    # Two checkpoints (recommended after Stage 2 training):
+    #   style_encoder from Stage 1, predictor_encoder from Stage 2
+    python scripts/extract_voicepack.py \
+        --model StyleTTS2/logs/kokoro_german/epoch_2nd_00001.pth \
+        --style-encoder-model StyleTTS2/logs/kokoro_german/epoch_1st_00002.pth \
+        --audio-dir dataset/audio/dm_daniel \
+        --output voices/dm_daniel.pt
+
+    # CPU (slower but works without GPU / while GPU is busy training)
+    python scripts/extract_voicepack.py \
+        --model StyleTTS2/logs/kokoro_german/epoch_1st_00002.pth \
+        --audio-dir dataset/audio/dm_daniel \
+        --output voices/dm_daniel.pt \
+        --device cpu
+
+Voicepack format: tensor of shape [510, 1, 256] (float32)
+  - 510 = max phoneme sequence length
+  - 256 = style_dim * 2 (128 for decoder/timbre + 128 for predictor/prosody)
+  - First 128 dims come from style_encoder (acoustic/timbre)
+  - Last 128 dims come from predictor_encoder (prosody)
+"""
+
+import argparse
+import math
+import sys
+from pathlib import Path
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.nn.utils.parametrizations import spectral_norm
+
+
+# ── StyleTTS2 model components (standalone, no external imports) ─────────
+
+
+class LearnedDownSample(nn.Module):
+    def __init__(self, layer_type, dim_in):
+        super().__init__()
+        self.layer_type = layer_type
+        if self.layer_type == "none":
+            self.conv = nn.Identity()
+        elif self.layer_type == "timepreserve":
+            self.conv = spectral_norm(
+                nn.Conv2d(
+                    dim_in,
+                    dim_in,
+                    kernel_size=(3, 1),
+                    stride=(2, 1),
+                    groups=dim_in,
+                    padding=(1, 0),
+                )
+            )
+        elif self.layer_type == "half":
+            self.conv = spectral_norm(
+                nn.Conv2d(
+                    dim_in,
+                    dim_in,
+                    kernel_size=(3, 3),
+                    stride=(2, 2),
+                    groups=dim_in,
+                    padding=1,
+                )
+            )
+        else:
+            raise RuntimeError(f"Unexpected downsample type: {self.layer_type}")
+
+    def forward(self, x):
+        return self.conv(x)
+
+
+class DownSample(nn.Module):
+    def __init__(self, layer_type):
+        super().__init__()
+        self.layer_type = layer_type
+
+    def forward(self, x):
+        if self.layer_type == "none":
+            return x
+        elif self.layer_type == "timepreserve":
+            return F.avg_pool2d(x, (2, 1))
+        elif self.layer_type == "half":
+            if x.shape[-1] % 2 != 0:
+                x = torch.cat([x, x[..., -1].unsqueeze(-1)], dim=-1)
+            return F.avg_pool2d(x, 2)
+        else:
+            raise RuntimeError(f"Unexpected downsample type: {self.layer_type}")
+
+
+class ResBlk(nn.Module):
+    def __init__(
+        self,
+        dim_in,
+        dim_out,
+        actv=nn.LeakyReLU(0.2),
+        normalize=False,
+        downsample="none",
+    ):
+        super().__init__()
+        self.actv = actv
+        self.normalize = normalize
+        self.downsample = DownSample(downsample)
+        self.downsample_res = LearnedDownSample(downsample, dim_in)
+        self.learned_sc = dim_in != dim_out
+        self._build_weights(dim_in, dim_out)
+
+    def _build_weights(self, dim_in, dim_out):
+        self.conv1 = spectral_norm(nn.Conv2d(dim_in, dim_in, 3, 1, 1))
+        self.conv2 = spectral_norm(nn.Conv2d(dim_in, dim_out, 3, 1, 1))
+        if self.normalize:
+            self.norm1 = nn.InstanceNorm2d(dim_in, affine=True)
+            self.norm2 = nn.InstanceNorm2d(dim_in, affine=True)
+        if self.learned_sc:
+            self.conv1x1 = spectral_norm(
+                nn.Conv2d(dim_in, dim_out, 1, 1, 0, bias=False)
+            )
+
+    def _shortcut(self, x):
+        if self.learned_sc:
+            x = self.conv1x1(x)
+        if self.downsample:
+            x = self.downsample(x)
+        return x
+
+    def _residual(self, x):
+        if self.normalize:
+            x = self.norm1(x)
+        x = self.actv(x)
+        x = self.conv1(x)
+        x = self.downsample_res(x)
+        if self.normalize:
+            x = self.norm2(x)
+        x = self.actv(x)
+        x = self.conv2(x)
+        return x
+
+    def forward(self, x):
+        x = self._shortcut(x) + self._residual(x)
+        return x / math.sqrt(2)
+
+
+class StyleEncoder(nn.Module):
+    def __init__(self, dim_in=48, style_dim=48, max_conv_dim=384):
+        super().__init__()
+        blocks = []
+        blocks += [spectral_norm(nn.Conv2d(1, dim_in, 3, 1, 1))]
+        repeat_num = 4
+        for _ in range(repeat_num):
+            dim_out = min(dim_in * 2, max_conv_dim)
+            blocks += [ResBlk(dim_in, dim_out, downsample="half")]
+            dim_in = dim_out
+        blocks += [nn.LeakyReLU(0.2)]
+        blocks += [spectral_norm(nn.Conv2d(dim_out, dim_out, 5, 1, 0))]
+        blocks += [nn.AdaptiveAvgPool2d(1)]
+        blocks += [nn.LeakyReLU(0.2)]
+        self.shared = nn.Sequential(*blocks)
+        self.unshared = nn.Linear(dim_out, style_dim)
+
+    def forward(self, x):
+        h = self.shared(x)
+        h = h.view(h.size(0), -1)
+        s = self.unshared(h)
+        return s
+
+
+# ── Main extraction logic ───────────────────────────────────────────────
+
+
+def extract_voicepack(
+    model_path: str,
+    audio_dir: str,
+    output_path: str,
+    num_samples: int = 200,
+    device: str = "auto",
+    style_encoder_model: str = None,
+):
+    import random
+
+    import soundfile as sf
+    import torchaudio
+
+    # ── Resolve device ───────────────────────────────────────────────────
+    if device == "auto":
+        if torch.cuda.is_available():
+            device = "cuda"
+        else:
+            device = "cpu"
+    print(f"Using device: {device}")
+
+    # ── Collect audio files ──────────────────────────────────────────────
+    audio_path = Path(audio_dir)
+    wav_files = sorted(audio_path.glob("*.wav"))
+    if not wav_files:
+        print(f"ERROR: No WAV files found in {audio_dir}")
+        sys.exit(1)
+
+    print(f"Found {len(wav_files):,} WAV files in {audio_dir}")
+
+    # Sample a subset for embedding
+    rng = random.Random(42)
+    if len(wav_files) > num_samples:
+        wav_files = rng.sample(wav_files, num_samples)
+    print(f"Using {len(wav_files)} samples for voicepack extraction")
+
+    # ── Build style encoders and load weights ────────────────────────────
+
+    def strip_prefix(state_dict, prefix="module."):
+        """Strip DataParallel 'module.' prefix from state dict keys if present."""
+        if any(k.startswith(prefix) for k in state_dict.keys()):
+            return {
+                k[len(prefix) :] if k.startswith(prefix) else k: v
+                for k, v in state_dict.items()
+            }
+        return state_dict
+
+    print(f"Loading checkpoint: {model_path}")
+    checkpoint = torch.load(model_path, map_location="cpu", weights_only=False)
+    net = checkpoint["net"]
+
+    # Kokoro-82M architecture: dim_in=64, style_dim=128, hidden_dim=512
+    style_encoder = StyleEncoder(dim_in=64, style_dim=128, max_conv_dim=512)
+    predictor_encoder = StyleEncoder(dim_in=64, style_dim=128, max_conv_dim=512)
+
+    # Load style_encoder weights — from a separate checkpoint if provided,
+    # otherwise from the main checkpoint. Stage 2 training can degrade the
+    # style_encoder (spectral_norm buffer drift), so loading from Stage 1
+    # is recommended when using a Stage 2 checkpoint.
+    if style_encoder_model:
+        print(f"Loading style_encoder from separate checkpoint: {style_encoder_model}")
+        se_checkpoint = torch.load(
+            style_encoder_model, map_location="cpu", weights_only=False
+        )
+        se_net = se_checkpoint["net"]
+        style_encoder.load_state_dict(strip_prefix(se_net["style_encoder"]))
+        se_epoch = se_checkpoint.get("epoch", "?")
+        print(f"  style_encoder loaded (epoch {se_epoch})")
+    else:
+        style_encoder.load_state_dict(strip_prefix(net["style_encoder"]))
+
+    # predictor_encoder is only trained in Stage 2 (train_second.py).
+    # After Stage 1 only, its weights are randomly initialized and produce
+    # exploded outputs even though parameter norms look reasonable.
+    # Detect this by checking output norm on a dummy input. If exploded,
+    # fall back to using style_encoder for both halves of the voicepack
+    # (same initialization Stage 2 uses — it deep-copies style_encoder).
+    predictor_encoder_trained = True
+    try:
+        predictor_encoder.load_state_dict(strip_prefix(net["predictor_encoder"]))
+        # Test output norm on a dummy mel-like input
+        with torch.no_grad():
+            dummy = torch.randn(1, 1, 80, 200)
+            test_out = predictor_encoder(dummy)
+            if test_out.norm().item() > 1e3:
+                predictor_encoder_trained = False
+    except Exception:
+        predictor_encoder_trained = False
+
+    if not predictor_encoder_trained:
+        print("  predictor_encoder appears untrained (Stage 1 only)")
+        print("  Using style_encoder for both acoustic and prosodic embeddings")
+        predictor_encoder.load_state_dict(
+            strip_prefix(se_net["style_encoder"])
+            if style_encoder_model
+            else strip_prefix(net["style_encoder"])
+        )
+
+    style_encoder = style_encoder.to(device).eval()
+    predictor_encoder = predictor_encoder.to(device).eval()
+
+    # Validate style encoder output norms
+    with torch.no_grad():
+        dummy = torch.randn(1, 1, 80, 200).to(device)
+        se_norm = style_encoder(dummy).norm().item()
+        pe_norm = predictor_encoder(dummy).norm().item()
+        print(f"  style_encoder output norm (dummy): {se_norm:.4f}")
+        print(f"  predictor_encoder output norm (dummy): {pe_norm:.4f}")
+        if se_norm < 0.5:
+            print(f"  WARNING: style_encoder norm is very low ({se_norm:.4f}).")
+            print(f"  This may indicate a collapsed encoder. Consider using")
+            print(f"  --style-encoder-model to load from a Stage 1 checkpoint.")
+
+    epoch = checkpoint.get("epoch", "?")
+    val_loss = checkpoint.get("val_loss", 0)
+    print(f"Loaded encoders from checkpoint")
+    print(f"  Epoch: {epoch}, Val loss: {val_loss:.4f}")
+
+    # ── Mel spectrogram transform ────────────────────────────────────────
+    # Must match StyleTTS2's training preprocessing exactly:
+    #   mel = to_mel(waveform)   # n_fft=2048, win=1200, hop=300, n_mels=80
+    #   mel = (log(1e-5 + mel) - mean) / std   # mean=-4, std=4
+    mel_transform = torchaudio.transforms.MelSpectrogram(
+        sample_rate=24000,
+        n_fft=2048,
+        win_length=1200,
+        hop_length=300,
+        n_mels=80,
+    ).to(device)
+
+    mel_mean = -4
+    mel_std = 4
+
+    # ── Extract style vectors ────────────────────────────────────────────
+    acoustic_styles = []
+    prosodic_styles = []
+    skipped = 0
+
+    print(f"\nExtracting style vectors...")
+    with torch.no_grad():
+        for i, wav_path in enumerate(wav_files):
+            # Load audio via soundfile (avoids torchcodec dependency)
+            try:
+                data, sr = sf.read(str(wav_path), dtype="float32")
+            except Exception as e:
+                print(f"  Skipping unreadable {wav_path}: {e}")
+                skipped += 1
+                continue
+            if data.ndim > 1:
+                data = data.mean(axis=1)  # mono
+            if sr != 24000:
+                # Resample using torchaudio
+                waveform = torch.from_numpy(data).unsqueeze(0)
+                waveform = torchaudio.functional.resample(waveform, sr, 24000)
+            else:
+                waveform = torch.from_numpy(data).unsqueeze(0)  # [1, samples]
+
+            waveform = waveform.to(device)
+
+            # Compute mel spectrogram (same as StyleTTS2's preprocess())
+            mel = mel_transform(waveform)  # [1, 80, T]
+            mel = (torch.log(1e-5 + mel) - mel_mean) / mel_std  # normalize
+
+            # Skip clips with very short mels (style encoder needs >= 80 frames)
+            if mel.shape[-1] < 80:
+                skipped += 1
+                continue
+
+            mel_input = mel.unsqueeze(1)  # [1, 1, 80, T] — add channel dim
+
+            # Run both style encoders
+            s_acoustic = style_encoder(mel_input)  # [1, 128]
+            s_prosodic = predictor_encoder(mel_input)  # [1, 128]
+
+            acoustic_styles.append(s_acoustic.cpu())
+            prosodic_styles.append(s_prosodic.cpu())
+
+            if (i + 1) % 50 == 0:
+                print(f"  Processed {i + 1}/{len(wav_files)} files...")
+
+    if skipped > 0:
+        print(f"  Skipped {skipped} files (too short)")
+
+    print(f"  Extracted {len(acoustic_styles)} style vectors")
+
+    if not acoustic_styles:
+        print("ERROR: No style vectors extracted. Check audio files.")
+        sys.exit(1)
+
+    # ── Average and build voicepack ──────────────────────────────────────
+    avg_acoustic = torch.cat(acoustic_styles, dim=0).mean(dim=0)  # [128]
+    avg_prosodic = torch.cat(prosodic_styles, dim=0).mean(dim=0)  # [128]
+
+    # Kokoro voicepack format: [510, 1, 256]
+    # First 128 dims = acoustic (decoder conditioning / timbre)
+    # Last 128 dims = prosodic (predictor conditioning / prosody)
+    combined = torch.cat([avg_acoustic, avg_prosodic], dim=0)  # [256]
+    voicepack = combined.unsqueeze(0).unsqueeze(0).expand(510, 1, 256).clone()
+
+    print(f"\nVoicepack shape: {tuple(voicepack.shape)}")
+    print(f"  Acoustic style norm: {avg_acoustic.norm():.4f}")
+    print(f"  Prosodic style norm: {avg_prosodic.norm():.4f}")
+
+    # ── Save ─────────────────────────────────────────────────────────────
+    output = Path(output_path)
+    output.parent.mkdir(parents=True, exist_ok=True)
+    torch.save(voicepack, str(output))
+    print(f"\nSaved voicepack: {output} ({output.stat().st_size / 1024:.1f} KB)")
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Extract voicepack from fine-tuned Kokoro model",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        epilog="""
+Examples:
+  # Stage 1 only (uses style_encoder for both halves)
+  python scripts/extract_voicepack.py \\
+      --model StyleTTS2/logs/kokoro_german/epoch_1st_00002.pth \\
+      --audio-dir dataset/audio/dm_daniel \\
+      --output voices/dm_daniel.pt
+
+  # After Stage 2 (recommended: style_encoder from Stage 1, predictor_encoder from Stage 2)
+  python scripts/extract_voicepack.py \\
+      --model StyleTTS2/logs/kokoro_german/epoch_2nd_00001.pth \\
+      --style-encoder-model StyleTTS2/logs/kokoro_german/epoch_1st_00002.pth \\
+      --audio-dir dataset/audio/dm_daniel \\
+      --output voices/dm_daniel.pt
+
+  # CPU (while GPU is busy training)
+  python scripts/extract_voicepack.py \\
+      --model StyleTTS2/logs/kokoro_german/epoch_1st_00002.pth \\
+      --device cpu
+""",
+    )
+    parser.add_argument(
+        "--model",
+        required=True,
+        help="Path to fine-tuned StyleTTS2 checkpoint (.pth). Used for predictor_encoder "
+        "(and style_encoder unless --style-encoder-model is specified).",
+    )
+    parser.add_argument(
+        "--style-encoder-model",
+        default=None,
+        help="Path to a separate checkpoint for loading the style_encoder weights. "
+        "Recommended: use a Stage 1 checkpoint here when --model is a Stage 2 "
+        "checkpoint, because Stage 2 can degrade the style_encoder.",
+    )
+    parser.add_argument(
+        "--audio-dir",
+        required=True,
+        help="Directory containing speaker's WAV files",
+    )
+    parser.add_argument(
+        "--output",
+        required=True,
+        help="Output voicepack path (.pt)",
+    )
+    parser.add_argument(
+        "--num-samples",
+        type=int,
+        default=200,
+        help="Number of audio samples to use for extraction (default: 200)",
+    )
+    parser.add_argument(
+        "--device",
+        default="auto",
+        choices=["auto", "cpu", "cuda"],
+        help="Device to run on: auto (default), cpu, or cuda",
+    )
+
+    args = parser.parse_args()
+    extract_voicepack(
+        model_path=args.model,
+        audio_dir=args.audio_dir,
+        output_path=args.output,
+        num_samples=args.num_samples,
+        device=args.device,
+        style_encoder_model=args.style_encoder_model,
+    )
+
+
+if __name__ == "__main__":
+    main()

training/prepare_arabic.py

@@ -0,0 +1,313 @@
+#!/usr/bin/env python3
+"""
+Kokoro Arabic (Palestinian) Training Data Preparation
+======================================================
+Builds StyleTTS2 train/val lists from two Palestinian Arabic datasets:
+
+  Speaker 0 "Kore"  : /home/ubuntu/data — Gemini TTS story segments
+                      (manifests/segments.jsonl + audio/*.wav + text/segments/*.txt)
+  Speaker 1 "Sirin" : hamdallah/Sofelia-Palestinian-TTS HF dataset
+                      (metadata.csv in LJSpeech format + wavs/*.wav)
+
+Both datasets are already 24kHz mono 16-bit WAV.
+
+G2P: misaki EspeakG2P(language='ar'), with fixups mapping phonemes that are
+not in Kokoro-82M's 178-token vocab onto unused-but-real vocab tokens so the
+phonemic contrasts of Arabic survive:
+
+  ħ (U+0127 voiceless pharyngeal fric.) -> ʰ (distinct h-like token)
+  ʕ (U+0295 voiced pharyngeal fric.)    -> ʁ (voiced uvular fric., unused by ar)
+  ˤ (U+02E4 pharyngealization/emphatic) -> ᵊ (distinct modifier token)
+  ̪ (U+032A dental diacritic)           -> removed (not phonemic)
+  [ ]                                   -> removed (espeak markup)
+
+Usage:
+    python3 scripts/prepare_arabic.py
+"""
+
+import json
+import random
+import sys
+from multiprocessing import Pool
+from pathlib import Path
+
+import soundfile as sf
+from tqdm import tqdm
+
+# ── Paths ────────────────────────────────────────────────────────────────────
+
+REPO = Path(__file__).resolve().parent.parent
+TRAINING_DIR = REPO / "training"
+KOKORO_CONFIG = TRAINING_DIR / "config.json"
+
+KORE_DIR = Path("/home/ubuntu/data")
+KORE_MANIFEST = KORE_DIR / "manifests" / "segments.jsonl"
+
+SOFELIA_DIR = Path(
+    "/home/ubuntu/.cache/huggingface/hub/datasets--hamdallah--Sofelia-Palestinian-TTS"
+    "/snapshots/316fd018d7855ada768639d40b0deb046be89558/dataset"
+)
+
+TRAIN_LIST = TRAINING_DIR / "train_list_ar.txt"
+VAL_LIST = TRAINING_DIR / "val_list_ar.txt"
+OOD_FILE = TRAINING_DIR / "OOD_texts_ar.txt"
+STATS_FILE = TRAINING_DIR / "stats_ar.json"
+
+# ── Filters / split ──────────────────────────────────────────────────────────
+
+MIN_DURATION_S = 2.0
+MAX_DURATION_S = 20.0  # caps worst-case batch memory (waveform discriminator)
+MIN_PHONEME_LEN = 10
+VAL_RATIO = 0.02
+N_OOD = 128  # entries reserved for OOD texts (excluded from training entirely)
+RANDOM_SEED = 42
+
+# ── Phoneme fixups (see module docstring) ────────────────────────────────────
+
+PHONEME_FIXUPS = {
+    "ħ": "ʰ",  # ħ -> ʰ
+    "ʕ": "ʁ",  # ʕ -> ʁ
+    "ˤ": "ᵊ",  # ˤ -> ᵊ
+    "̪": "",        # dental diacritic -> drop
+    "[": "",
+    "]": "",
+}
+
+_g2p = None
+_marbuta_cache = {}
+
+import json as _json
+import re
+
+_MARBUTA_WORD = re.compile(r"([ؠ-يـ]+ة)(?![ؠ-ي])")
+
+# Arabic punctuation → Latin equivalents that espeak passes through to the
+# phoneme string (the Arabic forms are silently dropped, so the model never
+# sees pause/question anchors without this).
+PUNCT_MAP = str.maketrans({"،": ",", "؛": ";", "؟": "?"})
+
+_LEXICON_PATH = REPO / "training" / "ar_lexicon.json"
+_lexicon = None
+_lexicon_re = None
+
+
+def _load_lexicon():
+    global _lexicon, _lexicon_re
+    if _lexicon is None:
+        if _LEXICON_PATH.exists():
+            entries = _json.loads(_LEXICON_PATH.read_text())
+            _lexicon = {k: v for k, v in entries.items() if not k.startswith("_")}
+        else:
+            _lexicon = {}
+        if _lexicon:
+            alts = "|".join(re.escape(w) for w in sorted(_lexicon, key=len, reverse=True))
+            _lexicon_re = re.compile(rf"(?<![ؠ-يـ])({alts})(?![ؠ-يـ])")
+    return _lexicon, _lexicon_re
+
+
+def normalize_arabic(text, g2p):
+    """Full text normalization: punctuation map + pronunciation lexicon + ة fix."""
+    text = text.translate(PUNCT_MAP)
+    lex, lex_re = _load_lexicon()
+    if lex_re is not None:
+        text = lex_re.sub(lambda m: lex[m.group(1)], text)
+    return normalize_taa_marbuta(text, g2p)
+
+
+def normalize_taa_marbuta(text, g2p):
+    """Fix espeak's OOV fallback that reads word-final ة as /t/ (e.g. قريبة→qriːbt).
+
+    espeak pronounces in-lexicon ة words correctly (مدرسة→a) but renders unknown
+    ones with a final t. For each ة-final word, G2P it alone: if the result ends
+    in 't', rewrite the trailing ة to a bare fatha (قريبَ), which espeak reads as
+    the correct pausal/dialect final 'a' (qriːba). Cached per word type.
+    """
+
+    def repl(m):
+        word = m.group(1)
+        fixed = _marbuta_cache.get(word)
+        if fixed is None:
+            try:
+                ph, _ = g2p(word)
+                broken = ph.strip().rstrip(".").endswith("t")
+            except Exception:
+                broken = False
+            fixed = word[:-1] + "َ" if broken else word  # ة → fatha
+            _marbuta_cache[word] = fixed
+        return fixed
+
+    return _MARBUTA_WORD.sub(repl, text)
+
+
+def _init_worker():
+    global _g2p
+    from misaki import espeak
+
+    _g2p = espeak.EspeakG2P(language="ar")
+
+
+def _phonemize(item):
+    """(idx, text) -> (idx, phonemes or None)"""
+    idx, text = item
+    try:
+        text = normalize_arabic(text, _g2p)
+        ph, _ = _g2p(text)
+    except Exception:
+        return idx, None
+    for old, new in PHONEME_FIXUPS.items():
+        ph = ph.replace(old, new)
+    return idx, ph.strip()
+
+
+def collect_kore():
+    """Speaker 0: Gemini story segments."""
+    entries = []
+    dropped = {"bad_duration": 0, "missing_file": 0, "empty_text": 0}
+    with open(KORE_MANIFEST) as f:
+        for line in f:
+            e = json.loads(line)
+            dur = e.get("durationSeconds", 0)
+            if not (MIN_DURATION_S <= dur <= MAX_DURATION_S):
+                dropped["bad_duration"] += 1
+                continue
+            seg_id = e["segmentId"]
+            wav = KORE_DIR / "audio" / f"{seg_id}.wav"
+            txt = KORE_DIR / "text" / "segments" / f"{seg_id}.txt"
+            if not wav.exists() or not txt.exists():
+                dropped["missing_file"] += 1
+                continue
+            text = txt.read_text().strip()
+            if not text:
+                dropped["empty_text"] += 1
+                continue
+            entries.append({"wav": str(wav), "text": text, "speaker": 0, "dur": dur})
+    print(f"Kore: kept {len(entries):,}  dropped {dropped}")
+    return entries
+
+
+def collect_sofelia():
+    """Speaker 1: Sofelia Palestinian TTS."""
+    entries = []
+    dropped = {"bad_duration": 0, "missing_file": 0, "empty_text": 0}
+    with open(SOFELIA_DIR / "metadata.csv") as f:
+        for line in f:
+            line = line.strip()
+            if not line:
+                continue
+            parts = line.split("|")
+            if len(parts) < 2:
+                continue
+            wav = SOFELIA_DIR / "wavs" / f"{parts[0]}.wav"
+            if not wav.exists():
+                dropped["missing_file"] += 1
+                continue
+            text = parts[1].strip()
+            if not text:
+                dropped["empty_text"] += 1
+                continue
+            info = sf.info(str(wav))
+            dur = info.frames / info.samplerate
+            if not (MIN_DURATION_S <= dur <= MAX_DURATION_S):
+                dropped["bad_duration"] += 1
+                continue
+            entries.append({"wav": str(wav), "text": text, "speaker": 1, "dur": dur})
+    print(f"Sofelia: kept {len(entries):,}  dropped {dropped}")
+    return entries
+
+
+def main():
+    vocab = json.load(open(KOKORO_CONFIG))["vocab"]
+
+    entries = collect_kore() + collect_sofelia()
+    print(f"Total entries before G2P: {len(entries):,}")
+
+    # ── G2P (parallel) ───────────────────────────────────────────────────
+    items = [(i, e["text"]) for i, e in enumerate(entries)]
+    with Pool(processes=8, initializer=_init_worker) as pool:
+        for idx, ph in tqdm(
+            pool.imap_unordered(_phonemize, items, chunksize=64),
+            total=len(items),
+            desc="G2P",
+        ):
+            entries[idx]["ipa"] = ph
+
+    # ── Validate against Kokoro vocab ────────────────────────────────────
+    valid = []
+    n_g2p_fail = n_short = n_unknown = 0
+    unknown_chars = {}
+    for e in entries:
+        ph = e.get("ipa")
+        if not ph:
+            n_g2p_fail += 1
+            continue
+        if len(ph) < MIN_PHONEME_LEN:
+            n_short += 1
+            continue
+        bad = [ch for ch in ph if ch not in vocab]
+        if bad:
+            n_unknown += 1
+            for ch in bad:
+                unknown_chars[ch] = unknown_chars.get(ch, 0) + 1
+            continue
+        valid.append(e)
+
+    print(f"Valid after G2P: {len(valid):,}")
+    print(f"  G2P failures   : {n_g2p_fail}")
+    print(f"  Too short      : {n_short}")
+    print(f"  Unknown phoneme: {n_unknown}")
+    if unknown_chars:
+        for ch, c in sorted(unknown_chars.items(), key=lambda x: -x[1])[:20]:
+            print(f"    {ch!r} U+{ord(ch):04X}: {c}")
+
+    if not valid:
+        print("ERROR: no valid entries")
+        sys.exit(1)
+
+    # ── Split: OOD / val / train ─────────────────────────────────────────
+    rng = random.Random(RANDOM_SEED)
+    rng.shuffle(valid)
+    ood_entries = valid[:N_OOD]
+    rest = valid[N_OOD:]
+    n_val = max(1, int(len(rest) * VAL_RATIO))
+    val_entries = rest[:n_val]
+    train_entries = rest[n_val:]
+
+    def write_list(path, ents):
+        with open(path, "w") as f:
+            for e in ents:
+                f.write(f"{e['wav']}|{e['ipa']}|{e['speaker']}\n")
+
+    write_list(TRAIN_LIST, train_entries)
+    write_list(VAL_LIST, val_entries)
+
+    with open(OOD_FILE, "w") as f:
+        f.write("\n".join(e["ipa"] for e in ood_entries) + "\n")
+
+    # ── Stats ────────────────────────────────────────────────────────────
+    def hours(ents):
+        return round(sum(e["dur"] for e in ents) / 3600, 2)
+
+    stats = {
+        "train": {"files": len(train_entries), "hours": hours(train_entries)},
+        "val": {"files": len(val_entries), "hours": hours(val_entries)},
+        "ood_texts": len(ood_entries),
+        "speakers": {
+            "0_kore": {
+                "files": sum(1 for e in valid if e["speaker"] == 0),
+                "hours": hours([e for e in valid if e["speaker"] == 0]),
+            },
+            "1_sirin": {
+                "files": sum(1 for e in valid if e["speaker"] == 1),
+                "hours": hours([e for e in valid if e["speaker"] == 1]),
+            },
+        },
+    }
+    with open(STATS_FILE, "w") as f:
+        json.dump(stats, f, indent=2)
+
+    print(json.dumps(stats, indent=2))
+    print(f"\nWrote {TRAIN_LIST}\n      {VAL_LIST}\n      {OOD_FILE}")
+
+
+if __name__ == "__main__":
+    main()

training/test_inference_arabic.py

@@ -0,0 +1,107 @@
+#!/usr/bin/env python3
+"""
+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  # noqa: E402
+from test_inference import convert_checkpoint  # noqa: E402
+
+TEST_SENTENCES = [
+    # Daily-life Palestinian dialect
+    "مرحبا، كيف حالك اليوم؟ إن شاء الله تكون مبسوط.",
+    "بدي أروح ع السوق أشتري خضرة وفواكه للبيت.",
+    # Pharyngeals and emphatics (ع، ح، ص، ض، ط)
+    "الصبح الساعة سبعة طلعت من البيت على الشغل.",
+    "حكيت مع عمي عن العرس اللي صار الأسبوع الماضي.",
+    # Numbers and prices
+    "هاد حقه خمسين شيكل بس، مش غالي.",
+    # Question prosody
+    "شو رأيك نطلع نتمشى شوي بعد العصر؟",
+    # Qaf-heavy (ق realization check)
+    "قلت لقاسم القهوة بالسوق قريبة من البنك القديم.",
+]
+
+
+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()