load_model.py

"""Load Nabra-7M-Distill, a 7.48M Arabic student distilled from Nabra-82M.

Use this rather than a plain `from kokoro import KModel`.

The config sets `hidden_channels` / `out_channels` on the decoder. Upstream Kokoro
hardcodes those at 1024/512, so stock `kokoro` raises TypeError on this config -- the decoder
is 53M of Kokoro's 82M and could not be shrunk without making the widths configurable. The
patched package is vendored here (76 KB of Python, defaults unchanged) so the repo is
self-contained and the 82M still loads through it untouched.

    from load_model import load
    model, pipeline, voice = load()
    audio = next(pipeline("مَرْحَبًا بِكُم", voice=voice))[2]
"""
import os, sys

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, os.path.join(HERE, "kokoro_patched"))   # must precede any kokoro import
sys.path.insert(0, HERE)

import torch
from arabic_g2p import EXTRA_SYMBOLS, clean_phonemes           # noqa: E402
from kokoro import KModel, KPipeline                            # noqa: E402
from kokoro import pipeline as _kp                              # noqa: E402

REPO = "oddadmix/Nabra-7M-Distill"


def load(device=None, weights="kokoro_arabic_7m.pth", config="config.json",
         voice="af_msa.pt"):
    device = device or ("cuda" if torch.cuda.is_available() else "cpu")
    m = KModel(repo_id=REPO, config=os.path.join(HERE, config),
               model=os.path.join(HERE, weights), disable_complex=True).eval()
    m.vocab.update(EXTRA_SYMBOLS)          # ʕ/ħ were trained on vocab gap slots 7/8
    m = m.to(device)
    _kp.LANG_CODES.setdefault("ar", "ar")
    p = KPipeline(lang_code="ar", repo_id=REPO, model=m)
    _orig = p.g2p
    p.g2p = lambda t: (clean_phonemes(_orig(t)[0]), _orig(t)[1])
    return m, p, torch.load(os.path.join(HERE, voice), map_location="cpu", weights_only=True)