Create app_quant.py

app_quant.py

@@ -0,0 +1,182 @@
+# ---------------------------------------------------------
+#   Nava Ultra-Fast CPU Inference (4-bit Quant + Caching)
+# ---------------------------------------------------------
+import gradio as gr
+import torch
+import soundfile as sf
+from pathlib import Path
+
+from transformers import (
+    AutoTokenizer,
+    AutoModelForCausalLM,
+    BitsAndBytesConfig
+)
+from peft import PeftModel
+from snac import SNAC
+
+# ---------------------------------------------------------
+# CONFIG
+# ---------------------------------------------------------
+MODEL_NAME = "rahul7star/nava1.0"
+LORA_NAME = "rahul7star/nava-audio"
+SNAC_MODEL_NAME = "rahul7star/nava-snac"
+
+SEQ_LEN = 240000
+TARGET_SR = 240000
+OUT_ROOT = Path("/tmp/data")
+OUT_ROOT.mkdir(exist_ok=True, parents=True)
+
+DEFAULT_TEXT = (
+    "राजनीतिज्ञों ने कहा कि उन्होंने निर्णायक मत को अनावश्यक रूप से "
+    "निर्धारित करने के लिए अफ़गान संविधान में काफी अस्पष्टता पाई थी"
+)
+
+DEVICE = "cpu"
+
+# ---------------------------------------------------------
+# QUANT CONFIG (4-BIT)
+# ---------------------------------------------------------
+quant_config = BitsAndBytesConfig(
+    load_in_4bit=True,
+    bnb_4bit_quant_type="nf4",
+    bnb_4bit_use_double_quant=True,
+    bnb_4bit_compute_dtype=torch.bfloat16,
+)
+
+# ---------------------------------------------------------
+# LOAD TOKENIZER (cached)
+# ---------------------------------------------------------
+print("🔄 Loading tokenizer...")
+tokenizer = AutoTokenizer.from_pretrained(
+    MODEL_NAME,
+    trust_remote_code=True
+)
+
+# ---------------------------------------------------------
+# LOAD BASE MODEL (4-bit CPU quant)
+# ---------------------------------------------------------
+print("🔄 Loading base model in 4-bit…")
+base_model = AutoModelForCausalLM.from_pretrained(
+    MODEL_NAME,
+    quantization_config=quant_config,
+    device_map={"": DEVICE},
+    torch_dtype=torch.bfloat16,
+    trust_remote_code=True
+)
+
+# ---------------------------------------------------------
+# LOAD LORA (merged on top)
+# ---------------------------------------------------------
+print("🔄 Loading LoRA weights…")
+model = PeftModel.from_pretrained(
+    base_model,
+    LORA_NAME,
+    device_map={"": DEVICE}
+).eval()
+
+# ---------------------------------------------------------
+# LOAD SNAC ONCE ONLY
+# ---------------------------------------------------------
+print("🔄 Loading SNAC…")
+snac_model = SNAC.from_pretrained(SNAC_MODEL_NAME).eval().to(DEVICE)
+
+
+# =========================================================
+#              INFERENCE FUNCTION
+# =========================================================
+def generate_audio_cpu_lora(text):
+
+    logs = []
+    logs.append("⚡ Running fast 4-bit CPU inference…")
+
+    # Tokens
+    soh = tokenizer.decode([128259])
+    eoh = tokenizer.decode([128260])
+    soa = tokenizer.decode([128261])
+    sos = tokenizer.decode([128257])
+    eot = tokenizer.decode([128009])
+    bos = tokenizer.bos_token
+
+    prompt = soh + bos + text + eot + eoh + soa + sos
+    inputs = tokenizer(prompt, return_tensors="pt").to(DEVICE)
+
+    # -----------------------------------------------------
+    # GENERATE SNAC TOKENS (FAST 4-bit)
+    # -----------------------------------------------------
+    with torch.inference_mode():
+        outputs = model.generate(
+            **inputs,
+            max_new_tokens=SEQ_LEN,
+            temperature=0.4,
+            top_p=0.9,
+            repetition_penalty=1.1,
+            do_sample=True,
+            eos_token_id=128258,
+            pad_token_id=tokenizer.pad_token_id
+        )
+
+    # Strip prompt
+    gen_ids = outputs[0, inputs['input_ids'].shape[1]:].tolist()
+
+    # Extract valid SNAC tokens
+    snac_min, snac_max = 128266, 156937
+    eos_id = 128258
+    eos_idx = gen_ids.index(eos_id) if eos_id in gen_ids else len(gen_ids)
+
+    snac_tokens = [t for t in gen_ids[:eos_idx] if snac_min <= t <= snac_max]
+
+    # -----------------------------------------------------
+    #  DECODE SNAC → AUDIO
+    # -----------------------------------------------------
+    l1, l2, l3 = [], [], []
+    frames = len(snac_tokens) // 7
+    snac_tokens = snac_tokens[:frames * 7]
+
+    for i in range(frames):
+        s = snac_tokens[i * 7:(i + 1) * 7]
+        l1.append((s[0] - snac_min) % 4096)
+        l2.extend([(s[1]-snac_min)%4096, (s[4]-snac_min)%4096])
+        l3.extend([(s[2]-snac_min)%4096, (s[3]-snac_min)%4096,
+                   (s[5]-snac_min)%4096, (s[6]-snac_min)%4096])
+
+    codes = [
+        torch.tensor(l1).unsqueeze(0),
+        torch.tensor(l2).unsqueeze(0),
+        torch.tensor(l3).unsqueeze(0)
+    ]
+
+    with torch.inference_mode():
+        z = snac_model.quantizer.from_codes(codes)
+        audio = snac_model.decoder(z)[0, 0].cpu().numpy()
+
+    # Remove crackles
+    if len(audio) > 2048:
+        audio = audio[2048:]
+
+    # Save WAV
+    out = OUT_ROOT / "tts_output_cpu_lora.wav"
+    sf.write(out, audio, TARGET_SR)
+
+    logs.append("🎧 Audio generated successfully")
+
+    return str(out), str(out), "\n".join(logs)
+
+
+# =========================================================
+#                   GRADIO UI
+# =========================================================
+with gr.Blocks() as demo:
+    gr.Markdown("## ⚡ Maya TTS — Ultra-Fast 4-bit CPU Inference")
+
+    txt = gr.Textbox(label="Enter text", value=DEFAULT_TEXT)
+    btn = gr.Button("Generate Audio")
+
+    audio = gr.Audio(label="Audio", type="filepath")
+    file = gr.File(label="Download")
+    logs = gr.Textbox(label="Logs")
+
+    btn.click(generate_audio_cpu_lora, [txt], [audio, file, logs])
+
+
+if __name__ == "__main__":
+    demo.launch()