app.py

"""
 
ZeroGPU (free tier) only attaches a GPU during execution of @spaces.GPU
functions. We:
  1. Install F5-TTS + Habibi-TTS at startup with --no-deps
  2. Load the model on CPU at startup (ZeroGPU has none during import)
  3. Move to CUDA + run inference inside a @spaces.GPU-decorated function
"""
import os
import sys
import subprocess
import tempfile
import traceback
from pathlib import Path


# ======================================================================
# 0. Install F5-TTS + Habibi-TTS at startup
# ======================================================================
def _ensure_tts_packages():
    try:
        import f5_tts      # noqa: F401
        import habibi_tts  # noqa: F401
        print("✅ f5_tts and habibi_tts already installed")
        return
    except ImportError:
        pass

    print("📦 Installing F5-TTS and Habibi-TTS with --no-deps...")
    for pkg_url in [
        "git+https://github.com/SWivid/F5-TTS.git",
        "git+https://github.com/SWivid/Habibi-TTS.git",
    ]:
        print(f"   → {pkg_url}")
        result = subprocess.run(
            [sys.executable, "-m", "pip", "install",
             "--no-cache-dir", "--no-deps", "-q", pkg_url],
            capture_output=True, text=True,
        )
        if result.returncode != 0:
            print(f"   ❌ install failed:\n{result.stderr}")
            raise RuntimeError(f"pip install failed for {pkg_url}")
        print(f"   ✅ installed")

_ensure_tts_packages()


# ======================================================================
# 1. Imports
# ======================================================================
import spaces                   # ← ZeroGPU decorator
import gradio as gr
import torch
import soundfile as sf
from huggingface_hub import hf_hub_download

from f5_tts.model import DiT
from f5_tts.infer.utils_infer import (
    load_model, load_vocoder, preprocess_ref_audio_text,
)
from habibi_tts.infer.utils_infer import infer_process


# ======================================================================
# 2. CONFIG
# ======================================================================
REPO_ID   = "NAMAA-Space/NAMAA-Saudi-TTS-V2"
CKPT_FILE = "model_2000.safetensors"
VOCAB     = "vocab.txt"


# ======================================================================
# 3. Load model on CPU at startup (no GPU available yet on ZeroGPU)
# ======================================================================
V1_BASE_CFG = dict(dim=1024, depth=22, heads=16,
                   ff_mult=2, text_dim=512, conv_layers=4)

print("📥 Downloading model weights + vocab from HF Hub...")
CKPT_PATH  = hf_hub_download(repo_id=REPO_ID, filename=CKPT_FILE)
VOCAB_PATH = hf_hub_download(repo_id=REPO_ID, filename=VOCAB)

print("🔧 Loading model on CPU (will move to GPU per-request on ZeroGPU)...")
# Load on CPU. Inside generate(), we move to cuda when GPU is attached.
MODEL = load_model(DiT, V1_BASE_CFG, CKPT_PATH,
                   vocab_file=VOCAB_PATH, device="cpu")
MODEL = MODEL.to(torch.float32).eval()
VOCODER = load_vocoder()
VOCODER = VOCODER.to("cpu")
print("✅ Model + vocoder ready on CPU")


# ======================================================================
# 4. Inference function (GPU-decorated for ZeroGPU)
# ======================================================================
@spaces.GPU(duration=60)   # request up to 60s of GPU per call
def generate(ref_audio, ref_text, gen_text, nfe_step, speed, remove_silence):
    if not ref_audio:
        raise gr.Error("Please upload a reference audio clip.")
    if not ref_text or not ref_text.strip():
        raise gr.Error("Please provide the reference transcript.")
    if not gen_text or not gen_text.strip():
        raise gr.Error("Please provide text to generate.")

    try:
        info = sf.info(ref_audio)
        dur = info.frames / info.samplerate
        if dur < 2.0:
            gr.Warning(f"Reference is only {dur:.1f}s — aim for 5-8s.")
        elif dur > 15.0:
            gr.Warning(f"Reference is {dur:.1f}s — will be truncated to 15s.")
    except Exception:
        pass

    try:
        # GPU is now available inside this decorated function.
        # Move model + vocoder onto cuda for this call.
        global MODEL, VOCODER
        device = "cuda" if torch.cuda.is_available() else "cpu"

        MODEL_GPU   = MODEL.to(device)
        VOCODER_GPU = VOCODER.to(device)

        ref_audio_p, ref_text_p = preprocess_ref_audio_text(
            ref_audio, ref_text.strip()
        )
        wave, sr, _ = infer_process(
            ref_audio_p, ref_text_p, gen_text.strip(),
            MODEL_GPU, VOCODER_GPU,
            nfe_step=int(nfe_step), speed=float(speed),
        )

        # Release the GPU copy (keep CPU copies in MODEL / VOCODER)
        if device == "cuda":
            del MODEL_GPU, VOCODER_GPU
            torch.cuda.empty_cache()

        if remove_silence:
            try:
                from f5_tts.infer.utils_infer import remove_silence_for_generated_wav
                with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp:
                    sf.write(tmp.name, wave, sr)
                    remove_silence_for_generated_wav(tmp.name)
                    wave, sr = sf.read(tmp.name)
                os.unlink(tmp.name)
            except Exception as e:
                gr.Warning(f"Silence removal skipped: {e}")

        duration = len(wave) / sr
        status = f"✅ Generated {duration:.1f}s audio at {sr}Hz on {device}"
        return (sr, wave), status

    except Exception as e:
        print(traceback.format_exc())
        raise gr.Error(f"Generation failed: {type(e).__name__}: {e}")


# ======================================================================
# 5. Examples
# ======================================================================
EXAMPLES = [
    ["examples/najdi_reference.wav",
     "تكفى طمني انا اليوم ماني بنايم",
     "مرحبا، كيف حالك اليوم؟", 32, 1.0, False],
    ["examples/najdi_reference.wav",
     "تكفى طمني انا اليوم ماني بنايم",
     "أهلا وسهلا بك في المملكة العربية السعودية", 32, 1.0, False],
    ["examples/najdi_reference.wav",
     "تكفى طمني انا اليوم ماني بنايم",
     "الرياض عاصمة المملكة العربية السعودية", 32, 1.0, False],
]


# ======================================================================
# 6. UI  (Gradio 6.x compatible)
# ======================================================================
DESCRIPTION = """
# 🇸🇦 Saudi TTS V2 — Voice Cloning TTS

Fine-tuned from [SWivid/Habibi-TTS](https://huggingface.co/SWivid/Habibi-TTS)
on ~18 hours of Najdi/Saudi Arabic audio. **Running on free ZeroGPU** —
first request per session takes ~30s to warm up.

**How to use:**
1. Upload a clean **5-8 second** reference clip (any Arabic voice)
2. Type the **exact transcript** of that clip
3. Type new Arabic text you want spoken in that voice
4. Click Generate
"""

ARTICLE = f"""
---
**Model:** [{REPO_ID}](https://huggingface.co/{REPO_ID}) •
**License:** CC-BY-NC-SA-4.0 •
**Architecture:** F5-TTS DiT (335M) + Vocos

Do not clone someone's voice without their consent.
"""


with gr.Blocks(title="Habibi-TTS Najdi") as demo:
    gr.Markdown(DESCRIPTION)

    with gr.Row():
        with gr.Column(scale=1):
            ref_audio_in = gr.Audio(
                label="Reference audio (5-8s, clean, single speaker)",
                type="filepath",
                sources=["upload", "microphone"],
            )
            ref_text_in = gr.Textbox(
                label="Reference transcript",
                placeholder="اكتب النص المنطوق في الصوت المرجعي",
                lines=2, rtl=True,
            )
            gen_text_in = gr.Textbox(
                label="Text to generate",
                placeholder="اكتب النص المطلوب",
                lines=3, rtl=True,
                value="مرحبا، كيف حالك اليوم؟",
            )
            with gr.Accordion("Advanced", open=False):
                nfe_step_in = gr.Slider(
                    minimum=8, maximum=64, value=32, step=4,
                    label="NFE steps (quality vs speed)",
                )
                speed_in = gr.Slider(
                    minimum=0.5, maximum=2.0, value=1.0, step=0.1,
                    label="Speech speed",
                )
                remove_silence_in = gr.Checkbox(
                    value=False,
                    label="Trim leading/trailing silence",
                )
            gen_btn = gr.Button("🎙 Generate", variant="primary", size="lg")

        with gr.Column(scale=1):
            audio_out = gr.Audio(
                label="Generated audio",
                type="numpy",
                autoplay=True,
            )
            status_out = gr.Textbox(
                label="Status", interactive=False, lines=2,
            )

    existing_examples = [ex for ex in EXAMPLES if os.path.exists(ex[0])]
    if existing_examples:
        gr.Examples(
            examples=existing_examples,
            inputs=[ref_audio_in, ref_text_in, gen_text_in,
                    nfe_step_in, speed_in, remove_silence_in],
            outputs=[audio_out, status_out],
            fn=generate,
            cache_examples=False,
            label="Examples (click to try)",
        )

    gr.Markdown(ARTICLE)

    gen_btn.click(
        fn=generate,
        inputs=[ref_audio_in, ref_text_in, gen_text_in,
                nfe_step_in, speed_in, remove_silence_in],
        outputs=[audio_out, status_out],
    )


if __name__ == "__main__":
    demo.queue(max_size=10).launch(theme=gr.themes.Soft())