# kikiri_pipeline — FlowTTS-style HF training pipeline for kikiri-tts

This package ports the **FlowTTS** training pipeline patterns into the kikiri-tts
fine-tuning workflow for **Kokoro-82M** (Urdu fine-tune on
[`humair025/koku`](https://huggingface.co/datasets/humair025/koku)).

It integrates:

  1. **HuggingFace dataset loading** — `datasets.load_dataset` with audio bytes
     decoding + 24 kHz mono resampling (Kokoro's format).
  2. **Urdu G2P phonemization** — `urdu_g2p` (vendored under `third_party/`)
     for OOD texts and fallback phonemization.
  3. **`KokoroTokenizer.encode`** — the single source of truth for converting
     IPA phoneme strings into Kokoro-82M token IDs (the 178-token vocab).
  4. **HuggingFace Hub push/pull** — auto-push of checkpoints, losses JSON,
     and TensorBoard logs to a HF model repo (with retry + pruning).
  5. **Resume from anywhere** — local path / HF Hub path / HF repo id / `'auto'`.
  6. **Stage 1 + Stage 2 training orchestrators** that wire all of the above
     into the patched StyleTTS2 training flow.

## Quick start

```bash
# 1. Install the pipeline deps (in addition to kikiri-tts's base deps)
pip install -r requirements-pipeline.txt

# 2. (Optional) Set HF token for upload + private dataset access
export HF_TOKEN=hf_xxx

# 3. Prepare the HF dataset → StyleTTS2 file lists + 24 kHz WAVs + token IDs
python scripts/prepare_dataset_hf.py \
    --config_path configs/config_urdu_ft.yml \
    --output_dir training \
    --max_rows 1000      # remove this for the full 30k-row dataset

# 4. Stage 1 training (acoustic reconstruction)
accelerate launch scripts/train_first.py \
    --config_path configs/config_urdu_ft.yml

# 5. Stage 2 training (prosody + adversarial)
accelerate launch scripts/train_second.py \
    --config_path configs/config_urdu_ft.yml

# 6. Resume Stage 2 from a HF Hub checkpoint
accelerate launch scripts/train_second.py \
    --config_path configs/config_urdu_ft.yml \
    --resume humair025/kikiri-urdu-v1/checkpoints/stage1/first_stage.pth
```

## What does the dataset look like?

[`humair025/koku`](https://huggingface.co/datasets/humair025/koku) is a 30k-row
Urdu TTS dataset with this schema:

| Column      | Type    | Notes                                          |
|-------------|---------|------------------------------------------------|
| `id`        | int64   | Row identifier                                 |
| `transcript`| string  | Full Urdu transcript                           |
| `voice`     | string  | Speaker label (e.g. `onyx`)                    |
| `text`      | string  | Same as `transcript` (alias)                   |
| `timestamp` | string  | Source timestamp                               |
| `audio`     | struct  | `{bytes: binary, path: string}` — HF Audio feat |
| `phonemes`  | string  | **Pre-computed IPA phonemes** (we use this!)   |

The pipeline **prefers the pre-computed `phonemes` column** for phonemization
(this is what the user means by "dataset already has col for phonemization").
`urdu_g2p` is only invoked for:
  * OOD warm-up sentences (`OOD_texts.txt` — phonemized on the fly)
  * Any rows where the `phonemes` column is empty / NaN
  * Inference-time text inputs (not part of this training pipeline)

## How does `KokoroTokenizer.encode` fit in?

The user explicitly asked for `kokoro.encode` from `KokoroTokenizer`. This is
the `KokoroTokenizer` class from
[`urdu_g2p_full/kokoro_compat/urdu_g2p_kokoro_tokenizer.py`](third_party/urdu_g2p/kokoro_compat/urdu_g2p_kokoro_tokenizer.py).

It:
  * Loads `kokoro_vocab.json` (the extended Kokoro vocab, 122 tokens)
  * Applies 3 transformations on the input IPA string:
      1. Multi-char phoneme rules: `t͡ʃ` → `ʧ` (id 133), `d͡ʒ` → `ʤ` (id 82)
      2. Decompose precomposed nasals: `ã` → `a` + `̃` (ids 43, 17), etc.
      3. Each remaining char maps directly to a Kokoro ID
  * Returns a `List[int]` of Kokoro token IDs

This is wired into the pipeline via `kikiri_pipeline.phonemizer.build_phonemizer()`:

```python
from kikiri_pipeline.phonemizer import build_phonemizer

ph = build_phonemizer()
ids = ph.encode_phonemes("paːkɪsˈt̪aːn")
#  -> [58, 43, 158, 53, 102, 61, 156, 62, 182, 43, 158, 56]
ids = ph.encode_text("پاکستان")  # full pipeline: text -> phonemes -> IDs
```

## Module layout

```
kikiri_pipeline/
  __init__.py          # Lazy-import hub for all submodules
  config.py            # YAML config loader with hf_pipeline block
  phonemizer.py        # UrduG2P + KokoroTokenizer integration
  audio_utils.py       # 24 kHz resampling, mel + F0 computation
  hf_dataset.py        # HF dataset loader (decode bytes, resample, encode)
  hf_hub.py            # Low-level push/pull helpers (file + folder + retry)
  checkpointer.py      # FlowTTS-style auto-push checkpointer for StyleTTS2
  resume.py            # Resolve --resume from local / HF / 'auto'
  train_logger.py      # JSON loss logger that pushes losses.json to HF
  data_prep.py         # HF dataset → StyleTTS2 train_list.txt + WAVs + token_ids/

scripts/
  prepare_dataset_hf.py   # CLI: HF dataset → StyleTTS2 training/ artefacts
  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.py       # (existing) Local MP3 → dataset pipeline (unchanged)
  prepare_training.py      # (existing) Local dataset prep (unchanged)
  extract_voicepack.py     # (existing) Voicepack extraction (unchanged)
  test_inference.py        # (existing) Inference test (unchanged)

configs/
  config_urdu_ft.yml       # Urdu fine-tuning config with hf_pipeline block
  config_german_ft.yml     # (existing) German config

third_party/
  urdu_g2p/                # Vendored urdu_g2p package (self-contained)
    urdu_g2p/              # Main G2P module
    kokoro_compat/         # KokoroTokenizer + kokoro_vocab.json

training/
  kokoro_symbols.py        # (existing) Kokoro 178-token symbol mapping
  config.json              # (existing) Kokoro model config
  OOD_texts.txt            # (existing, but regenerated by data_prep)

requirements-pipeline.txt # Extra deps for the HF pipeline
PIPELINE_README.md         # This file
```

## Config: the `hf_pipeline:` block

Add this block to any training config YAML to enable the pipeline. All
keys are optional with strong defaults — see
[`configs/config_urdu_ft.yml`](configs/config_urdu_ft.yml) for a full example:

```yaml
hf_pipeline:
  # ---- Dataset ----
  dataset_id:        "humair025/koku"
  dataset_split:     "train"
  dataset_audio_col: "audio"
  dataset_text_col:  "text"
  dataset_phonemes_col: "phonemes"     # pre-computed IPA column
  dataset_voice_col: "voice"
  sample_rate:       24000              # Kokoro = 24 kHz mono
  streaming:         true               # set false if you have >10 GB disk
  val_ratio:         0.05
  min_phoneme_len:   5
  max_phoneme_len:   510                # PLBERT max_position_embeddings - 2

  # ---- HF Hub push ----
  hf_repo:           "humair025/kikiri-urdu-v1"
  hf_token_env:      "HF_TOKEN"
  upload_every:      1                   # push ckpt every N epochs
  keep_last_n_local: 2
  upload_tensorboard: true
  tb_log_upload_freq: 200
  losses_upload_freq: 50
  hf_dry_run:        false

  # ---- Resume ----
  resume:            null                # local path / HF path / 'auto'
```

Override any of these via CLI flags:

```bash
python scripts/train_first.py \
    --config_path configs/config_urdu_ft.yml \
    --hf_repo humair025/kikiri-urdu-v1 \
    --upload_every 5 \
    --hf_dry_run        # save locally but don't push to HF
```

## Training stages

### Stage 1 (`train_first.py`)

Acoustic reconstruction: trains the text encoder + decoder + PLBERT to
reconstruct mel spectrograms. Reads `epochs_1st` from the config.
Saves `first_stage.pth` at the end.

```bash
accelerate launch scripts/train_first.py \
    --config_path configs/config_urdu_ft.yml
```

### Stage 2 (`train_second.py`)

Prosody prediction + adversarial training: trains the duration + F0 +
energy predictor on top of the Stage 1 acoustic model. Activates GAN +
WavLM SLM adversarial loss from `joint_epoch` (default 3). Reads
`epochs_2nd` from the config. Saves `second_stage.pth` at the end.

```bash
accelerate launch scripts/train_second.py \
    --config_path configs/config_urdu_ft.yml
```

### Standalone smoke-test mode

Both `train_first.py` and `train_second.py` support `--standalone` for
smoke-testing the HF pipeline WITHOUT needing the StyleTTS2 submodule
checked out. This runs a minimal training loop that exercises:

  * HF dataset loading
  * 24 kHz audio decode + resample
  * KokoroTokenizer.encode for phonemes → token IDs
  * Checkpointer save + (optional) HF Hub push
  * LossLogger JSON logging + (optional) HF Hub push
  * Resume from local checkpoint

```bash
# Stage 1 smoke test (10 steps, no HF upload)
python scripts/train_first.py \
    --config_path configs/config_urdu_ft.yml \
    --standalone --max_steps 10 --hf_dry_run

# Stage 2 smoke test (10 steps, no HF upload)
python scripts/train_second.py \
    --config_path configs/config_urdu_ft.yml \
    --standalone --max_steps 10 --hf_dry_run
```

**Note**: standalone mode does NOT implement full StyleTTS2 semantics
(PLBERT, monotonic alignment, ISTFTNet decoder, multi-period discriminator,
WavLM SLM loss, etc.). For real training, use delegate mode with the
StyleTTS2 submodule checked out:

```bash
git submodule update --init --recursive
```

## Checkpoint + loss layout

After training, you'll have:

```
logs/kokoro-urdu/
  stage1/
    epoch_stage1_00000.pth       # StyleTTS2-format checkpoint
    epoch_stage1_00001.pth       # (pruned to keep_last_n_local)
    ...
    final.pth                    # final Stage 1 model
    first_stage.pth              # alias for final.pth (kikiri convention)
    losses.json                  # full loss history
    config_snapshot.json
    run_info.json
    runs/                        # TensorBoard logs (if enabled)
  stage2/
    epoch_stage2_00000.pth
    ...
    final.pth
    second_stage.pth
    losses.json
    config_snapshot.json
    run_info.json
    runs/
```

Each `.pth` checkpoint contains:

```python
{
  'net': {
    'bert':          state_dict,
    'bert_encoder':  state_dict,
    'predictor':     state_dict,
    'text_encoder':  state_dict,
    'decoder':       state_dict,
  },
  'epoch':     int,
  'iters':     int,
  'stage':     'stage1' | 'stage2',
  'losses':    {'loss_mel': ..., 'loss_gen': ..., ...},
  'val_loss':  float,     # if computed
  'optimizers': {...},   # if passed to checkpointer.save()
  'schedulers': {...},   # if passed
  'extra':     {...},
}
```

If `hf_repo` is set and `hf_dry_run` is `false`, every checkpoint is also
pushed to:

```
https://huggingface.co/{hf_repo}/blob/main/checkpoints/{stage_tag}/epoch_{stage_tag}_NNNNN.pth
https://huggingface.co/{hf_repo}/blob/main/checkpoints/{stage_tag}/final.pth
https://huggingface.co/{hf_repo}/blob/main/logs/{stage_tag}/losses.json
https://huggingface.co/{hf_repo}/tree/main/tensorboard_logs/{stage_tag}
```

## Resume

Resume is supported from any of:

```bash
# 1. Local path
--resume logs/kokoro-urdu/stage1/final.pth

# 2. HF Hub file path
--resume humair025/kikiri-urdu-v1/checkpoints/stage1/final.pth

# 3. HF repo id (auto-find latest)
--resume humair025/kikiri-urdu-v1

# 4. 'auto' (look local first, then HF)
--resume auto
```

The resume helper is in [`kikiri_pipeline/resume.py`](kikiri_pipeline/resume.py).
It mirrors FlowTTS's `resume.py` but also recognises kikiri/StyleTTS2-style
filenames (`epoch_2nd_00009.pth`, `first_stage.pth`, `second_stage.pth`).

## Listing checkpoints in a HF repo

```bash
python scripts/train_first.py --list_hub_checkpoints humair025/kikiri-urdu-v1
```

This prints all numbered checkpoints in the repo, sorted by step/epoch.

## API reference (quick)

```python
from kikiri_pipeline.config import load_pipeline_config
from kikiri_pipeline.phonemizer import build_phonemizer
from kikiri_pipeline.hf_dataset import HFDatasetLoader
from kikiri_pipeline.checkpointer import KikiriCheckpointer
from kikiri_pipeline.train_logger import LossLogger
from kikiri_pipeline.resume import resolve_resume_path

# 1. Load config
cfg = load_pipeline_config("configs/config_urdu_ft.yml")

# 2. Build phonemizer (urdu_g2p + KokoroTokenizer)
ph = build_phonemizer(warn_only=True)
ids = ph.encode_phonemes("paːkɪsˈt̪aːn")  # -> [58, 43, 158, ...]

# 3. Build HF dataset loader
loader = HFDatasetLoader(
    dataset_id=cfg.hf.dataset_id,
    split=cfg.hf.dataset_split,
    audio_col=cfg.hf.dataset_audio_col,
    phonemes_col=cfg.hf.dataset_phonemes_col,
    sample_rate=cfg.hf.sample_rate,
    token_encoder=ph.encode_phonemes,
    streaming=cfg.hf.streaming,
)
for row in loader:
    # row.audio: np.float32 @ 24 kHz mono
    # row.token_ids: List[int] (Kokoro IDs)
    # row.phonemes: str (IPA)
    # row.speaker: str
    ...

# 4. Build loss logger + checkpointer
loss_logger = LossLogger(
    log_dir="logs/kokoro-urdu/stage1",
    stage="stage1",
    hf_repo=cfg.hf.hf_repo,
    hf_token=cfg.hf.resolve_token(),
    losses_upload_freq=cfg.hf.losses_upload_freq,
    dry_run=cfg.hf.hf_dry_run,
)
checkpointer = KikiriCheckpointer.from_config(
    cfg, stage_tag="stage1", loss_logger=loss_logger)

# 5. Inside training loop
checkpointer.save(
    net={"bert": ..., "decoder": ..., ...},
    epoch=epoch,
    iters=step,
    losses={"loss_mel": mel_loss, "loss_gen": gen_loss, ...},
    val_loss=val_loss,
)
loss_logger.log(step=step, epoch=epoch, losses={"loss_mel": ...})

# 6. End of training
checkpointer.final_upload(
    net={"bert": ..., "decoder": ..., ...},
    epoch=epoch, iters=step, losses=final_losses,
)

# 7. Resume
ckpt_path = resolve_resume_path(
    resume_spec="humair025/kikiri-urdu-v1",
    checkpoint_dir="logs/kokoro-urdu/stage1",
    hf_repo=cfg.hf.hf_repo,
    hf_token=cfg.hf.resolve_token(),
    subdir="checkpoints/stage1",
)
```

## What this is NOT

  * **Not a replacement for StyleTTS2.** The patched StyleTTS2 submodule
    is still required for real training. `train_first.py` / `train_second.py`
    in `scripts/` delegate to `StyleTTS2/train_first.py` /
    `StyleTTS2/train_second.py` (subprocess).
  * **Not a Kokoro inference package.** For inference, use the `kokoro/`
    submodule (KModel + KPipeline) or the upstream `hexgrad/kokoro` package.
  * **Not a phonemizer.** `urdu_g2p` is the phonemizer. We integrate it,
    we don't replace it.

## See also

  * [`docs/TRAINING_GUIDE.md`](docs/TRAINING_GUIDE.md) — kikiri-tts training guide
  * [`docs/ARCHITECTURE.md`](docs/ARCHITECTURE.md) — Kokoro-82M architecture notes
  * [`configs/config_urdu_ft.yml`](configs/config_urdu_ft.yml) — Urdu FT config
  * [`requirements-pipeline.txt`](requirements-pipeline.txt) — extra deps