Update app.py

app.py

@@ -1,16 +1,91 @@
-import gradio as gr
-from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+# Somali TTS • Hugging Face Space
+# Author: <Your‑Name>
+# Model  : Somalitts/xuseen  (Orpheus‑style SNAC)
 
-model_id = "Somalitts/XUSEEN"  # ama magaca saxda ah ee modelkaaga
+import os, torch, gradio as gr, torchaudio.transforms as T
+from unsloth import FastLanguageModel
+from snac import SNAC
 
-tokenizer = AutoTokenizer.from_pretrained(model_id)
-model = AutoModelForCausalLM.from_pretrained(model_id, device_map="auto")
-pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
+# ------------------------------------------------------------------
+# 0. Environment ‑ for HF Spaces (better memory use)
+os.environ["BITSANDBYTES_NOWELCOME"] = "1"
+device = "cuda" if torch.cuda.is_available() else "cpu"
 
-def generate(prompt):
-    output = pipe(prompt, max_new_tokens=100, do_sample=True, temperature=0.7)
-    return output[0]["generated_text"]
+# ------------------------------------------------------------------
+# 1. Load LLM (4‑bit) + Tokenizer
+MODEL_ID = "Somalitts/XUSEEN"     # <- beddel haddii loo baahdo
+model, tokenizer = FastLanguageModel.from_pretrained(
+    model_name      = MODEL_ID,
+    load_in_4bit    = True,
+    max_seq_length  = 4096,
+)
+FastLanguageModel.for_inference(model)         # 2× speed
 
-demo = gr.Interface(fn=generate, inputs="text", outputs="text", title="Somali LLM")
+# ------------------------------------------------------------------
+# 2. Load SNAC decoder (24 kHz)
+snac = SNAC.from_pretrained("hubertsiuzdak/snac_24khz").to(device)
 
-demo.launch()
+# Special‑tokens (waa ku dhex jira Somalitts models)
+AUDIO_START = 128266
+END_TEXT    = 128009
+SOH, EOH    = 128259, 128260
+SOS, EOS    = 128257, 128258      # AUDIO_START‑128256+{1,2}
+
+@torch.inference_mode()
+def tts_somali(text: str):
+    """Generate speech waveform (24 kHz, numpy) from Somali text."""
+    if not text.strip():
+        return 24_000, None
+
+    # ───── Build prompt ─────
+    prompt_ids = tokenizer.encode(text, add_special_tokens=True) + [END_TEXT]
+    input_ids  = torch.tensor([[SOH] + prompt_ids + [EOH]], dtype=torch.int64).to(device)
+
+    # ───── Generate audio codes ─────
+    gen = model.generate(
+        input_ids       = input_ids,
+        max_new_tokens  = 1200,
+        do_sample       = True,
+        temperature     = 0.7,
+        top_p           = 0.95,
+        eos_token_id    = EOS,
+    )[0]
+
+    # ───── Extract 7‑code tuples & distribute to 3 layers ─────
+    s_idx = (gen == SOS).nonzero(as_tuple=True)[0] + 1
+    e_idx = (gen == EOS).nonzero(as_tuple=True)[0]
+    codes = (gen[s_idx:e_idx] - AUDIO_START).tolist()
+    codes = codes[: (len(codes) // 7) * 7]           # multiple‑of‑7
+
+    l1, l2, l3 = [], [], []
+    for i in range(len(codes) // 7):
+        c = codes[7 * i : 7 * i + 7]
+        l1.append(c[0])
+        l2.append(c[1] - 4096)
+        l2.append(c[4] - 16384)
+        l3.extend([c[2] - 8192, c[3] - 12288, c[5] - 20480, c[6] - 24576])
+
+    codes_torch = [torch.tensor(l).unsqueeze(0).to(device) for l in (l1, l2, l3)]
+    audio = snac.decode(codes_torch)
+    return 24_000, audio.squeeze().cpu().numpy()
+
+# ------------------------------------------------------------------
+# 3. Gradio UI
+DESCRIPTION = """
+<b>Somali Text‑to‑Speech</b><br>
+Ku qor qoraal Af‑Soomaali ah oo gaaban (<i>≤ 200 eray</i>) si aad u maqasho cod nadiif ah 24 kHz.
+"""
+
+demo = gr.Interface(
+    fn          = tts_somali,
+    inputs      = gr.Textbox(lines=4, label="Qoraalkaaga halkan geli"),
+    outputs     = gr.Audio(type="numpy", label="Codka la soo saar(24 kHz)"),
+    title       = "Somali TTS – Xuseen",
+    description = DESCRIPTION,
+    theme       = "default",
+    examples    = ["Soomaaliya waxa ay leedahay xeebta ugu dheer Afrika.",
+                   "Magaca caasimadda Soomaaliya waa Muqdisho."],
+)
+
+if __name__ == "__main__":
+    demo.launch()