# kikiri-tts Urdu Pipeline — Quick Start

## What's in this zip

```
kikiri-tts-main/
├── kikiri_pipeline/           # FlowTTS-style HF pipeline (12 modules, ~3K LOC)
│   ├── __init__.py
│   ├── config.py              # YAML config loader with hf_pipeline block
│   ├── phonemizer.py          # UrduG2P + KokoroTokenizer178 integration
│   ├── kokoro_tokenizer_178.py  # 178-token Kokoro vocab (IDs in [0,177])
│   ├── kokoro_weights.py      # Download + convert original Kokoro-82M weights
│   ├── real_kokoro_model.py   # Wraps 81.76M Kokoro weights as nn.Parameters
│   ├── audio_utils.py         # 24 kHz resampling, mel + F0 computation
│   ├── hf_dataset.py          # HF dataset loader (decode bytes, resample, encode)
│   ├── hf_hub.py              # Push/pull checkpoints + losses + TB logs to HF
│   ├── checkpointer.py       # FlowTTS-style auto-push for StyleTTS2 checkpoints
│   ├── resume.py              # Resume from local / HF Hub / 'auto'
│   ├── train_logger.py        # JSON loss logger that pushes to HF
│   └── data_prep.py           # HF dataset → StyleTTS2 train_list.txt + WAVs
│
├── scripts/
│   ├── train_first.py         # Stage 1 training (delegates to StyleTTS2/train_first.py)
│   ├── train_second.py        # Stage 2 training (delegates to StyleTTS2/train_second.py)
│   ├── prepare_dataset_hf.py  # CLI: HF dataset → StyleTTS2 training artefacts
│   ├── prepare_dataset.py     # (existing) Local MP3 → dataset pipeline
│   ├── prepare_training.py    # (existing) Local dataset prep
│   ├── extract_voicepack.py   # (existing) Voicepack extraction
│   └── test_inference.py       # (existing) Inference test
│
├── configs/
│   ├── config_urdu_ft.yml     # Urdu FT config with hf_pipeline block
│   └── config_german_ft.yml   # (existing) German config
│
├── StyleTTS2/                 # semidark/StyleTTS2 at commit b1956da
│   ├── train_first.py         # Stage 1 training (patched for CPU + Kokoro)
│   ├── train_second.py        # Stage 2 training (patched)
│   ├── meldataset.py          # Dataset loader (uses kokoro_symbols)
│   ├── models.py              # StyleTTS2 model definitions
│   ├── losses.py              # Loss functions (WavLM SLM, STFT, GAN)
│   ├── kokoro_symbols.py      # 178-token Kokoro vocab
│   ├── kokoro_tb_utils.py     # Kokoro TensorBoard inference utilities
│   ├── monotonic_align_src/   # resemble-ai/monotonic_align source (build via setup.py)
│   └── Utils/                 # Pretrained utility models (ASR, JDC, PLBERT)
│       ├── ASR/               # ASR alignment model (94 MB)
│       ├── JDC/               # F0 pitch extractor (21 MB)
│       └── PLBERT/            # Phoneme-level BERT (25 MB)
│
├── third_party/
│   └── urdu_g2p/              # Vendored Urdu G2P package (self-contained)
│       ├── urdu_g2p/          # Main G2P module
│       └── kokoro_compat/     # KokoroTokenizer + kokoro_vocab.json
│
├── training/                  # Training data + Kokoro symbols
│   ├── kokoro_symbols.py      # 178-token symbol mapping
│   ├── config.json            # Kokoro-82M config (downloaded by prepare_base_weights)
│   ├── train_list.txt         # 4 sample rows (Urdu)
│   ├── val_list.txt           # 1 sample row
│   ├── OOD_texts.txt          # 20 Urdu OOD sentences (phonemized)
│   ├── speaker_map.json       # {"onyx": 0}
│   ├── audio/onyx/*.wav       # 5 sample WAVs (24 kHz mono)
│   └── token_ids/*.json       # Per-row Kokoro token IDs
│
├── PIPELINE_README.md         # Full pipeline documentation
├── requirements-pipeline.txt  # Extra deps (datasets, librosa, tensorboard, etc.)
├── pyproject.toml             # Base kikiri-tts deps
└── docs/                      # Training guide, troubleshooting, architecture
```

## Quick start

```bash
# 1. Unzip
unzip kikiri-tts-urdu-pipeline.zip
cd kikiri-tts-main

# 2. Install dependencies
pip install -r requirements-pipeline.txt
cd StyleTTS2/monotonic_align_src && pip install . && cd ../..

# 3. Set HF token (for uploads)
export HF_TOKEN=hf_xxx

# 4. Download + convert original Kokoro-82M weights (327 MB)
python -c "
from kikiri_pipeline.kokoro_weights import prepare_base_weights
prepare_base_weights('training/kokoro_base.pth', hf_token='$HF_TOKEN')
"

# 5. Prepare dataset (5 rows for smoke test, remove --max_rows for full 30k)
python scripts/prepare_dataset_hf.py \
    --config_path configs/config_urdu_ft.yml \
    --output_dir training \
    --max_rows 100

# 6. Stage 1 training (real StyleTTS2)
cd StyleTTS2
accelerate launch train_first.py -p ../configs/config_urdu_ft.yml
cd ..

# 7. Stage 2 training
cd StyleTTS2
accelerate launch train_second.py -p ../configs/config_urdu_ft.yml
cd ..

# 8. Or use the standalone smoke test (no StyleTTS2 submodule needed):
python scripts/train_first.py \
    --config_path configs/config_urdu_ft.yml \
    --standalone --max_steps 10 \
    --hf_repo your-username/kikiri-urdu-v1 \
    --upload_every 5 --upload_tensorboard
```

## What was verified end-to-end

- ✅ Original Kokoro-82M weights downloaded + converted to StyleTTS2 format (81.76M params)
- ✅ 178-token KokoroTokenizer (all IDs in [0, 177])
- ✅ HF dataset loading from humair025/koku (24 kHz mono resampling)
- ✅ UrduG2P + KokoroTokenizer.encode integration
- ✅ StyleTTS2 dataloader reads our train_list.txt correctly
- ✅ StyleTTS2 model loads the converted Kokoro weights (all 5 components)
- ✅ Real training runs forward + backward through actual Kokoro weights
- ✅ Checkpoints + TensorBoard logs + losses.json pushed to HF Hub
- ✅ Resume from local + HF Hub checkpoints

## HF Hub repo (example)

https://huggingface.co/humair025/kikiri-urdu-v1

Contains:
- `checkpoints/stage1/{epoch_stage1_00000.pth, final.pth}` (327 MB each)
- `tensorboard_logs/stage1/events.out.tfevents.*` (scalars + histograms)
- `logs/stage1/losses.json` (per-step loss history)
- `kikiri_pipeline/*.py` (full pipeline source)
- `configs/config_urdu_ft.yml`
- `dataset_prep/{train_list.txt, val_list.txt, OOD_texts.txt}`
- `PIPELINE_README.md`

## Notes

- The `kokoro_base.pth` (327 MB) is NOT in the zip — regenerate it via step 4 above.
- The StyleTTS2 code has been patched: hardcoded `"cuda"` → `accelerator.device` and `F0_curve.device` so it runs on CPU too.
- For real GPU training, install the CUDA version of torch: `pip install torch torchaudio --index-url https://download.pytorch.org/whl/cu121`