"""Load oddadmix/Nabra-Saudi-7M, a 7.48M Saudi-dialect Arabic voice.

Use this rather than a plain `from kokoro import KModel`: this config sets the decoder's
`hidden_channels` / `out_channels`, which upstream Kokoro hardcodes, so the stock package raises
a TypeError on it. The patched package is vendored here with its defaults unchanged.

    from load_model import load
    model, pipeline, voice = load()
    audio = next(pipeline("يا أهلاً وسهلاً، كيف حالكم؟", voice=voice))[2]   # 24 kHz

Saudi dialect text is normally written WITHOUT full tashkeel, and that is what this model
expects. Adding MSA diacritics makes it worse, not better: an MSA diacritizer imposes case
endings a dialect speaker does not pronounce.
"""
import os, sys, torch

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, os.path.join(HERE, "kokoro_patched"))
sys.path.insert(0, HERE)

from kokoro import KModel, KPipeline          # noqa: E402
from arabic_g2p import ArabicG2P, EXTRA_SYMBOLS, clean_phonemes   # noqa: E402
from kokoro import pipeline as kpipeline_mod  # noqa: E402

kpipeline_mod.LANG_CODES.setdefault("ar", "ar")
REPO = "oddadmix/Nabra-Saudi-7M"


def load(device=None, weights="nabra_saudi_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().to(device)
    m.vocab.update(EXTRA_SYMBOLS)
    p = KPipeline(lang_code="ar", repo_id=REPO, model=m)
    g = ArabicG2P(diacritize=False)           # dialect: use the text as written
    p.g2p = lambda t: (clean_phonemes(g.process(t.replace("\u0640", ""))[1]), t)
    return m, p, torch.load(os.path.join(HERE, voice), map_location="cpu", weights_only=True)