app.py

"""
KasbahTTS V0 — Algerian Dardja text-to-speech.

Zero-shot voice cloning demo for MenaVoice/KasbahTTS-V0, an F5-TTS (DiT flow
matching) model fine-tuned on Algerian Dardja.

`spaces` must be imported before torch so ZeroGPU can patch CUDA init.
"""

import os
import random
import re
import tempfile

try:
    import spaces

    USING_ZEROGPU = True
except ImportError:
    USING_ZEROGPU = False

import gradio as gr
import soundfile as sf
import torch
from f5_tts.infer.utils_infer import (
    load_model,
    load_vocoder,
    preprocess_ref_audio_text,
    remove_silence_for_generated_wav,
)
from f5_tts.model import DiT
from huggingface_hub import hf_hub_download

from habibi_tts.infer.utils_infer import infer_process


def gpu_decorator(func):
    # Only wrap with the ZeroGPU allocator when actually running on ZeroGPU
    # hardware (env var set by HF). On CPU/paid-GPU spaces this is a no-op.
    if USING_ZEROGPU and os.environ.get("SPACES_ZERO_GPU"):
        return spaces.GPU(duration=120)(func)
    return func


MODEL_REPO = "MenaVoice/KasbahTTS-V0"
CKPT_FILE = "ALGERIA.safetensors"
V1_CFG = dict(dim=1024, depth=22, heads=16, ff_mult=2, text_dim=512, conv_layers=4)

# On ZeroGPU, CUDA exists only *inside* a @spaces.GPU call — never at import time.
# So load everything on CPU here and move to GPU inside generate().
ON_ZEROGPU = USING_ZEROGPU and bool(os.environ.get("SPACES_ZERO_GPU"))
GPU_DEVICE = "cuda" if (ON_ZEROGPU or torch.cuda.is_available()) else "cpu"

ckpt_path = hf_hub_download(MODEL_REPO, CKPT_FILE)
vocab_path = hf_hub_download(MODEL_REPO, "vocab.txt")

vocoder = load_vocoder(vocoder_name="vocos", is_local=False, device="cpu")
model = load_model(DiT, V1_CFG, ckpt_path, vocab_file=vocab_path, device="cpu")

EXAMPLES_DIR = os.path.join(os.path.dirname(__file__), "examples")
# (reference wav, exact transcript of that wav) — transcripts must match the
# audio verbatim; all refs are kept under f5-tts's 12s reference clip limit.
EXAMPLE_REFS = [
    (
        os.path.join(EXAMPLES_DIR, "1.wav"),
        "مَثَلاً ما تَلْڨاشْ واحدْ يقول لك شْجْرة قْلَقتني، شْجْرة مَرَضَتْلي حْيَاتي. تْلْڨاه غِيرْ يْڨُول لك هَذ الإنْسَان فْلَان وْ فْلَان وْ فْلَان. تسما مَشَاكِلْنا كَامَل مَنْ عَنْد النَّاس.",
    ),
    (
        os.path.join(EXAMPLES_DIR, "2.wav"),
        "الطالبات يروحوا ليها يتغداو، كنت أنا نروح بعد لي كور باش نتغدى، وكان كاين واحد الشاب دائماً يشوف فيا ويحب يحكي معايا، ويعاملني معاملة خاصة.",
    ),
    (
        os.path.join(EXAMPLES_DIR, "3.wav"),
        "الدراهم ولا القدره المعيشيه كيما الفلاح عمر هذا كان عادي معندوش و مخصوش عايش على قد حال",
    ),
]

# The model was trained on unvocalized Arabic script only: Latin letters and
# digits fall outside its vocabulary and produce unpredictable audio.
UNSUPPORTED = re.compile(r"[A-Za-z0-9٠-٩]")


@gpu_decorator
def generate(
    ref_audio_orig,
    ref_text,
    gen_text,
    speed,
    nfe_step,
    cfg_strength,
    cross_fade,
    remove_sil,
    seed,
    progress=gr.Progress(),
):
    if not ref_audio_orig:
        raise gr.Error("Please provide a reference audio clip (5-15 seconds works best).")
    if not gen_text.strip():
        raise gr.Error("Please enter some Algerian Dardja text to synthesize.")

    if UNSUPPORTED.search(gen_text):
        gr.Warning(
            "Latin letters and digits are outside the model's vocabulary. "
            "Spell numbers out in Arabic words (ثلاثة, not 3)."
        )

    if seed is None or int(seed) < 0:
        seed = random.randint(0, 2**31 - 1)
    seed = int(seed)
    torch.manual_seed(seed)

    # Leaving ref_text blank triggers Whisper auto-transcription inside f5-tts.
    ref_audio, ref_text = preprocess_ref_audio_text(ref_audio_orig, ref_text or "", show_info=gr.Info)

    # Move to GPU now that we're inside the @spaces.GPU context (loaded on CPU at startup).
    model.to(GPU_DEVICE)
    vocoder.to(GPU_DEVICE)
    try:
        # KasbahTTS is a *specialized* single-dialect checkpoint, so no dialect
        # token is prepended — infer_process defaults dialect_id to None.
        wave, sample_rate, _ = infer_process(
            ref_audio,
            ref_text,
            gen_text,
            model,
            vocoder,
            cross_fade_duration=cross_fade,
            nfe_step=int(nfe_step),
            cfg_strength=cfg_strength,
            speed=speed,
            show_info=gr.Info,
            progress=progress,
            device=GPU_DEVICE,
        )
    finally:
        # ZeroGPU tears down CUDA after the call — return the model to CPU so the
        # next call starts from a valid copy.
        if ON_ZEROGPU:
            model.to("cpu")
            vocoder.to("cpu")
            torch.cuda.empty_cache()

    if remove_sil:
        with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as f:
            tmp_path = f.name
        try:
            sf.write(tmp_path, wave, sample_rate)
            remove_silence_for_generated_wav(tmp_path)
            wave, sample_rate = sf.read(tmp_path)
        finally:
            os.unlink(tmp_path)

    return (sample_rate, wave), ref_text, seed


with gr.Blocks(title="KasbahTTS V0", theme=gr.themes.Soft()) as demo:
    gr.Markdown(
        """
        # 🇩🇿 KasbahTTS V0 — Algerian Dardja TTS

        Zero-shot voice cloning for **Algerian Dardja** (الدارجة الجزائرية).
        Upload a few seconds of any voice, type Dardja text, and hear it spoken back in that voice.

        Model: [`MenaVoice/KasbahTTS-V0`](https://huggingface.co/MenaVoice/KasbahTTS-V0) ·
        F5-TTS architecture, fine-tuned from Habibi-TTS · MIT licensed.
        """
    )

    with gr.Row():
        with gr.Column():
            ref_audio = gr.Audio(label="Reference voice", type="filepath")
            ref_text = gr.Textbox(
                label="Reference transcript (optional)",
                lines=2,
                placeholder="Leave blank to auto-transcribe with Whisper…",
            )
            gen_text = gr.Textbox(
                label="Text to generate (Arabic script only)",
                lines=4,
                placeholder="واش راك خويا، لاباس عليك؟",
                rtl=True,
            )
            run = gr.Button("🎤 Generate", variant="primary", size="lg")

        with gr.Column():
            audio_out = gr.Audio(label="Generated speech", interactive=False)
            ref_text_out = gr.Textbox(label="Reference transcript used", interactive=False, lines=2)
            seed_out = gr.Number(label="Seed used", interactive=False, precision=0)

            with gr.Accordion("Advanced settings", open=False):
                speed = gr.Slider(0.5, 1.5, value=1.0, step=0.05, label="Speed")
                nfe_step = gr.Slider(
                    8, 64, value=32, step=8, label="NFE steps",
                    info="More steps = higher quality, slower.",
                )
                cfg_strength = gr.Slider(0.5, 5.0, value=2.0, step=0.1, label="CFG strength")
                cross_fade = gr.Slider(0.0, 0.5, value=0.15, step=0.05, label="Cross-fade (s)")
                remove_sil = gr.Checkbox(label="Trim long silences", value=False)
                seed = gr.Number(label="Seed (-1 = random)", value=-1, precision=0)

    inputs = [ref_audio, ref_text, gen_text, speed, nfe_step, cfg_strength, cross_fade, remove_sil, seed]
    outputs = [audio_out, ref_text_out, seed_out]

    run.click(generate, inputs=inputs, outputs=outputs)

    available_refs = [(wav, text) for wav, text in EXAMPLE_REFS if os.path.exists(wav)]
    if available_refs:
        gen_samples = [
            "واش راك خويا، لاباس عليك؟ صباح الخير.",
            "من القصبة للعالم، هذا أول موديل يهدر بالدارجة الجزائرية.",
            "اليوم الجو شباب، قلت نخرج نتمشى شوية في وسط البلاد.",
        ]
        gr.Examples(
            examples=[
                [wav, text, gen_samples[i % len(gen_samples)]]
                for i, (wav, text) in enumerate(available_refs)
            ],
            inputs=[ref_audio, ref_text, gen_text],
        )

    gr.Markdown(
        """
        ---
        ### Known limitations

        - **Arabic script only.** No French code-switching — `ça va` will not work.
        - **No digits.** Write `ثلاثة`, not `3`.
        - **No diacritics (تشكيل).** Use plain, unvocalized Arabic.
        - **Occasional repetition.** Try `nfe_step` 64, adjust `cfg_strength`, or split long text.

        Built with ❤️ for Algeria by MenaVoice — *من القصبة للعالم*
        """
    )


if __name__ == "__main__":
    demo.queue().launch()