app imp v 2

app_quant.py

@@ -1,376 +1,207 @@
-# app.py
-"""
-Optimized Maya + LoRA + SNAC inference app
-Loads model + LoRA + SNAC at startup, vectorized SNAC unpacking,
-torch.compile for speed, and Gradio UI with presets & emotions.
-"""
-
-import os
-import time
-import traceback
-from pathlib import Path
-
+# app_quant_superfast.py
 import gradio as gr
-import numpy as np
-import soundfile as sf
 import torch
+import soundfile as sf
+from pathlib import Path
+import traceback
+import time
 
 from transformers import AutoTokenizer, AutoModelForCausalLM
-from peft import PeftModel, LoraConfig
+from peft import PeftModel
 from snac import SNAC
 
-# -------------------------
-# Config - update if needed
-# -------------------------
+
+# =========================================================
+# CONFIG
+# =========================================================
 MODEL_NAME = "rahul7star/nava1.0"
 LORA_NAME = "rahul7star/nava-audio"
-SNAC_MODEL_NAME = "hubertsiuzdak/snac_24khz"  # or your snac model name
+SNAC_MODEL_NAME = "rahul7star/nava-snac"
 TARGET_SR = 24000
+
 OUT_ROOT = Path("/tmp/data")
 OUT_ROOT.mkdir(exist_ok=True, parents=True)
 
-# default text
-DEFAULT_TEXT = ("राजनीतिज्ञों ने कहा कि उन्होंने निर्णायक मत को अनावश्यक रूप से "
-                "निर्धारित करने के लिए अफ़गान संविधान में काफी अस्पष्टता पाई थी")
-
-# Preset characters
-PRESET_CHARACTERS = {
-    "Male American": {
-        "description": "Realistic male voice in the 20s age with an american accent. High pitch, raspy timbre, brisk pacing, neutral tone delivery at medium intensity, viral_content domain, short_form_narrator role, neutral delivery",
-        "example_text": DEFAULT_TEXT
-    },
-    "Female British": {
-        "description": "Realistic female voice in the 30s age with a british accent. Normal pitch, throaty timbre, conversational pacing, sarcastic tone delivery at low intensity, podcast domain, interviewer role, formal delivery",
-        "example_text": DEFAULT_TEXT
-    },
-    "Robot": {
-        "description": "Creative, ai_machine_voice character. Male voice in their 30s with an american accent. High pitch, robotic timbre, slow pacing, sad tone at medium intensity.",
-        "example_text": DEFAULT_TEXT
-    },
-    "Singer": {
-        "description": "Creative, animated_cartoon character. Male voice in their 30s with an american accent. High pitch, deep timbre, slow pacing, sarcastic tone at medium intensity.",
-        "example_text": DEFAULT_TEXT
-    },
-    "Custom (use LoRA)": {
-        "description": "Use your trained LoRA voice. Edit description or use preset emotions.",
-        "example_text": DEFAULT_TEXT
-    },
-}
-
-EMOTION_TAGS = [
-    "<neutral>", "<angry>", "<chuckle>", "<cry>", "<disappointed>", "<excited>",
-    "<gasp>", "<giggle>", "<laugh>", "<laugh_harder>", "<sarcastic>", "<sigh>",
-    "<sing>", "<whisper>"
-]
-
-# Generation sequence length config
-SEQ_LEN_CPU = 4096
-SEQ_LEN_GPU = 240000
-MAX_NEW_TOKENS_CPU = 1024
-MAX_NEW_TOKENS_GPU = 240000
-
-# -------------------------
-# Detect device & bitsandbytes
-# -------------------------
+DEFAULT_TEXT = "राजनीतिज्ञों ने कहा कि उन्होंने निर्णायक मत को अनावश्यक रूप से निर्धारित करने के लिए अफ़गान संविधान में काफी अस्पष्टता पाई थी"
+
 HAS_CUDA = torch.cuda.is_available()
 DEVICE = "cuda" if HAS_CUDA else "cpu"
 
-bnb_available = False
-if HAS_CUDA:
-    try:
-        from transformers import BitsAndBytesConfig
-        bnb_available = True
-    except Exception:
-        bnb_available = False
+MAX_NEW = 240000 if HAS_CUDA else 1024
 
-print(f"[init] Device: {DEVICE} | CUDA: {HAS_CUDA} | bnb: {bnb_available}")
 
-# -------------------------
-# Load tokenizer
-# -------------------------
-print("[init] Loading tokenizer...")
+print(f"[INIT] Running on: {DEVICE}")
+
+
+# =========================================================
+# LOAD TOKENIZER (light)
+# =========================================================
+print("[INIT] Loading tokenizer...")
 tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True)
-# safe tokenizer tweaks
-tokenizer.model_max_length = SEQ_LEN_GPU if HAS_CUDA else SEQ_LEN_CPU
-tokenizer.padding_side = "left"
-
-# Cache the special tokens strings once
-SOH_TOK = tokenizer.decode([128259])
-EOH_TOK = tokenizer.decode([128260])
-SOA_TOK = tokenizer.decode([128261])
-SOS_TOK = tokenizer.decode([128257])
-EOT_TOK = tokenizer.decode([128009])
-BOS_TOK = tokenizer.bos_token
-
-# -------------------------
-# Load base model + LoRA once (optimized)
-# -------------------------
-print("[init] Loading base model and attaching LoRA... this may take some time")
-model = None
-try:
-    if HAS_CUDA and bnb_available:
-        # GPU + bnb (4-bit) path
-        quant_config = BitsAndBytesConfig(
-            load_in_4bit=True,
-            bnb_4bit_quant_type="nf4",
-            bnb_4bit_use_double_quant=True,
-            bnb_4bit_compute_dtype=torch.bfloat16
-        )
-        base_model = AutoModelForCausalLM.from_pretrained(
-            MODEL_NAME,
-            quantization_config=quant_config,
-            device_map="auto",
-            trust_remote_code=True,
-        )
-        # attach LoRA
-        model = PeftModel.from_pretrained(base_model, LORA_NAME, device_map="auto")
-        SEQ_LEN = SEQ_LEN_GPU
-        MAX_NEW_TOKENS = MAX_NEW_TOKENS_GPU
-        print("[init] Base+LoRA loaded on GPU (4-bit).")
-    else:
-        # CPU (or GPU without bnb): load base model with low_cpu_mem_usage then attach PEFT
-        base_model = AutoModelForCausalLM.from_pretrained(
-            MODEL_NAME,
-            torch_dtype=torch.float32,
-            device_map={"": "cpu"},
-            low_cpu_mem_usage=True,
-            trust_remote_code=True,
-        )
-        model = PeftModel.from_pretrained(base_model, LORA_NAME, device_map={"": "cpu"})
-        SEQ_LEN = SEQ_LEN_CPU
-        MAX_NEW_TOKENS = MAX_NEW_TOKENS_CPU
-        print("[init] Base+LoRA loaded on CPU (FP32).")
-except Exception as e:
-    print("[init] Error loading model:", e)
-    traceback.print_exc()
-    raise
-
-# Speed/quality configs
-try:
-    # prefer use_cache and disable extra outputs
-    model.config.use_cache = True
-    model.config.output_attentions = False
-    model.config.output_hidden_states = False
-    # torch.compile if available (PyTorch >= 2.0)
-    try:
-        print("[init] Trying torch.compile for model (may speed up inference)...")
-        # choose a safe compile mode
-        if DEVICE == "cuda":
-            model = torch.compile(model)
-        else:
-            # CPU safe mode: reduce overhead
-            model = torch.compile(model, mode="reduce-overhead", fullgraph=False)
-        print("[init] torch.compile applied (if supported).")
-    except Exception as e:
-        print("[init] torch.compile not applied:", e)
-except Exception:
-    pass
 
-# If GPU, set half precision where possible
-if DEVICE == "cuda":
-    try:
-        model.half()
-        print("[init] model.half() applied for GPU.")
-    except Exception:
-        pass
+
+# =========================================================
+# LOAD BASE MODEL + LORA (ONCE)
+# =========================================================
+print("[INIT] Loading base model...")
+if HAS_CUDA:
+    # GPU 4-bit
+    from transformers import BitsAndBytesConfig
+    quant = BitsAndBytesConfig(
+        load_in_4bit=True,
+        bnb_4bit_compute_dtype=torch.bfloat16,
+        bnb_4bit_use_double_quant=True,
+        bnb_4bit_quant_type="nf4",
+    )
+    base_model = AutoModelForCausalLM.from_pretrained(
+        MODEL_NAME,
+        quantization_config=quant,
+        device_map="auto",
+        trust_remote_code=True,
+    )
+else:
+    # CPU 8-bit (MUCH faster than fp32)
+    base_model = AutoModelForCausalLM.from_pretrained(
+        MODEL_NAME,
+        device_map={"": "cpu"},
+        trust_remote_code=True,
+        load_in_8bit=True,
+    )
+
+print("[INIT] Attaching LoRA...")
+model = PeftModel.from_pretrained(base_model, LORA_NAME, device_map={"": DEVICE})
+
+# 🔥 Merge LoRA weights permanently = big speedup
+print("[INIT] Merging LoRA -> base weights...")
+model = model.merge_and_unload()
 
 model.eval()
+torch.set_grad_enabled(False)
+
+print("[INIT] Model ready.")
 
-# -------------------------
-# Load SNAC Model (once)
-# -------------------------
-print("[init] Loading SNAC decoder...")
+
+# =========================================================
+# LOAD SNAC DECODER
+# =========================================================
+print("[INIT] Loading SNAC...")
 snac_model = SNAC.from_pretrained(SNAC_MODEL_NAME).eval().to(DEVICE)
-if DEVICE == "cuda":
-    try:
-        snac_model.half()
-    except Exception:
-        pass
-print("[init] SNAC loaded.")
-
-# -------------------------
-# Utility: Vectorized SNAC unpacking
-# -------------------------
-def snac_tokens_to_levels(snac_tokens):
-    """
-    snac_tokens: list[int] (already filtered to SNAC-range and multiples of 7)
-    returns tuple(torch.Tensor l1, l2, l3) each on DEVICE
-    """
-    if len(snac_tokens) == 0:
-        return None, None, None
-
-    # ensure length multiple of 7
-    frames = len(snac_tokens) // 7
-    total = frames * 7
-    arr = torch.tensor(snac_tokens[:total], dtype=torch.long, device=DEVICE)
-    frames_tensor = arr.view(frames, 7)
-
-    snac_min = 128266
-    # l1: (frames,)
-    l1 = ((frames_tensor[:, 0] - snac_min) % 4096).unsqueeze(0)
-    # l2: interleave frames[:,1] and frames[:,4]
-    l2_col1 = ((frames_tensor[:, 1] - snac_min) % 4096)
-    l2_col2 = ((frames_tensor[:, 4] - snac_min) % 4096)
-    l2 = torch.stack([l2_col1, l2_col2], dim=1).reshape(1, -1)
-    # l3: frames[:,2], frames[:,3], frames[:,5], frames[:,6]
-    l3_cols = torch.stack([
-        (frames_tensor[:, 2] - snac_min) % 4096,
-        (frames_tensor[:, 3] - snac_min) % 4096,
-        (frames_tensor[:, 5] - snac_min) % 4096,
-        (frames_tensor[:, 6] - snac_min) % 4096,
-    ], dim=1)
-    l3 = l3_cols.reshape(1, -1)
-
-    return l1, l2, l3
-
-# -------------------------
-# Inference function (single unified, uses loaded model + snac_model)
-# -------------------------
-def generate_audio_with_lora(text: str,
-                             description: str = "",
-                             emotion_tag: str = "<neutral>",
-                             max_new_tokens_override: int = None):
-    """
-    Build prompt using description+emotion (if provided), run model.generate,
-    extract SNAC tokens, decode via snac_model -> return filepath + logs
-    """
+
+print("[INIT COMPLETE]")
+
+
+# =========================================================
+# PRE-COMPUTED TOKENS (SPEED-UP)
+# =========================================================
+soh = tokenizer.decode([128259])
+eoh = tokenizer.decode([128260])
+soa = tokenizer.decode([128261])
+sos = tokenizer.decode([128257])
+eot = tokenizer.decode([128009])
+bos = tokenizer.bos_token
+snac_min, snac_max = 128266, 156937
+eos_id = 128258
+
+
+# =========================================================
+# FAST INFERENCE
+# =========================================================
+def generate_audio(text):
     logs = []
     t0 = time.time()
     try:
-        # build prompt: keep consistent with Maya prompt format
-        if description:
-            prompt_text = f'<description="{description}"> {text}'
-        else:
-            prompt_text = text
-
-        # prepend special tokens (cached)
-        prompt = SOH_TOK + BOS_TOK + prompt_text + EOT_TOK + EOH_TOK + SOA_TOK + SOS_TOK
-
-        logs.append(f"[INFO] DEVICE={DEVICE} | prompt_len={len(prompt)}")
+        logs.append(f"[INFO] Using {DEVICE}")
 
-        # Tokenize
-        inputs = tokenizer(prompt, return_tensors="pt", padding=True).to(DEVICE)
+        prompt = f"{soh}{bos}{text}{eot}{eoh}{soa}{sos}"
 
-        # choose max_new
-        max_new = MAX_NEW_TOKENS
-        if max_new_tokens_override:
-            max_new = min(max_new, int(max_new_tokens_override))
-        if DEVICE == "cpu":
-            max_new = min(max_new, MAX_NEW_TOKENS_CPU)
+        inputs = tokenizer(prompt, return_tensors="pt").to(DEVICE)
 
-        # Generate
         with torch.inference_mode():
-            outputs = model.generate(
+            output = model.generate(
                 **inputs,
-                max_new_tokens=max_new,
-                temperature=0.4,
+                max_new_tokens=MAX_NEW,
+                do_sample=True,
+                temperature=0.5,
                 top_p=0.9,
                 repetition_penalty=1.1,
-                do_sample=True,
-                eos_token_id=128258,
+                eos_token_id=eos_id,
                 pad_token_id=tokenizer.pad_token_id,
+                use_cache=True,          # 🔥 faster
             )
 
-        gen_ids = outputs[0, inputs["input_ids"].shape[1]:].tolist()
-        logs.append(f"[INFO] Generated {len(gen_ids)} tokens")
+        gen = output[0, inputs['input_ids'].shape[1]:].tolist()
 
-        # filter 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]
+        logs.append(f"[INFO] Generated token count: {len(gen)}")
 
-        # make frames multiple of 7
+        # strip non-SNAC
+        if eos_id in gen:
+            gen = gen[:gen.index(eos_id)]
+
+        snac_tokens = [t for t in gen if snac_min <= t <= snac_max]
         frames = len(snac_tokens) // 7
-        if frames == 0:
-            logs.append("[WARN] No SNAC frames found in generated tokens.")
-            return None, None, "\n".join(logs)
+        snac_tokens = snac_tokens[:frames * 7]
 
-        l1, l2, l3 = snac_tokens_to_levels(snac_tokens)
-        if l1 is None:
-            logs.append("[ERROR] SNAC conversion failed.")
+        if frames == 0:
+            logs.append("[WARN] No SNAC frames!")
             return None, None, "\n".join(logs)
 
-        # decode via SNAC (keep on DEVICE)
+        # unpack
+        l1 = []
+        l2 = []
+        l3 = []
+
+        for i in range(frames):
+            s = snac_tokens[i*7:(i+1)*7]
+            l1.append((s[0]-snac_min) % 4096)
+            l2 += [(s[1]-snac_min)%4096, (s[4]-snac_min)%4096]
+            l3 += [(s[2]-snac_min)%4096, (s[3]-snac_min)%4096,
+                   (s[5]-snac_min)%4096, (s[6]-snac_min)%4096]
+
+        codes = [
+            torch.tensor(l1, device=DEVICE).unsqueeze(0),
+            torch.tensor(l2, device=DEVICE).unsqueeze(0),
+            torch.tensor(l3, device=DEVICE).unsqueeze(0),
+        ]
+
+        # decode → audio
         with torch.inference_mode():
-            z_q = snac_model.quantizer.from_codes([l1.to(DEVICE), l2.to(DEVICE), l3.to(DEVICE)])
-            audio_tensor = snac_model.decoder(z_q)[0, 0]
-            # ensure float32 on CPU before saving
-            audio = audio_tensor.cpu().float().numpy()
+            zq = snac_model.quantizer.from_codes(codes)
+            audio = snac_model.decoder(zq)[0, 0].cpu().numpy()
+
+        # trim
+        audio = audio[2048:]
 
-        # trim warmup samples
-        if len(audio) > 2048:
-            audio = audio[2048:]
+        out = OUT_ROOT / "tts.wav"
+        sf.write(out, audio, TARGET_SR)
 
-        out_path = OUT_ROOT / "tts_output_lora.wav"
-        sf.write(str(out_path), audio, TARGET_SR)
-        logs.append(f"[OK] Audio saved: {out_path} (duration {len(audio)/TARGET_SR:.2f}s)")
+        logs.append(f"[OK] Audio saved {out}")
+        logs.append(f"[TIME] {time.time() - t0:.2f}s")
 
-        logs.append(f"[TIME] Elapsed: {time.time()-t0:.2f}s")
-        return str(out_path), str(out_path), "\n".join(logs)
+        return str(out), str(out), "\n".join(logs)
 
     except Exception as e:
-        tb = traceback.format_exc()
-        logs.append(f"[ERROR] {e}\n{tb}")
+        logs.append(str(e))
+        logs.append(traceback.format_exc())
         return None, None, "\n".join(logs)
 
-# -------------------------
-# Small wrapper to provide default description/emotion insertion
-# -------------------------
-def generate_for_ui(text: str, preset: str, description: str, emotion: str):
-    # choose description: preset overrides empty
-    if preset and preset in PRESET_CHARACTERS:
-        # if user selected "Custom (use LoRA)" keep user-provided description
-        if preset != "Custom (use LoRA)":
-            desc = PRESET_CHARACTERS[preset]["description"]
-        else:
-            desc = description or PRESET_CHARACTERS[preset]["description"]
-    else:
-        desc = description
-
-    # include emotion tag appended to description to make model aware
-    combined_desc = f"{desc} Emotion: {emotion}"
-    audio_fp, dl_fp, logs = generate_audio_with_lora(text, description=combined_desc, emotion_tag=emotion)
-    return audio_fp, dl_fp, logs
-
-# -------------------------
-# Gradio UI
-# -------------------------
-with gr.Blocks(title="NAVA — Maya1 + LoRA + SNAC (Optimized)", css=".gradio-container {max-width: 1400px}") as demo:
-    gr.Markdown("# 🪶 NAVA — Maya1 + LoRA + SNAC (Optimized)\nGenerate emotional Hindi speech using Maya1 base + your LoRA adapter.")
-    with gr.Row():
-        with gr.Column(scale=3):
-            gr.Markdown("## Inference (CPU/GPU auto)\nType text + pick a preset or write description manually.")
-            text_in = gr.Textbox(label="Enter Hindi text", value=DEFAULT_TEXT, lines=3)
-            preset_select = gr.Dropdown(label="Select Preset Character", choices=list(PRESET_CHARACTERS.keys()), value="Male American")
-            description_box = gr.Textbox(label="Voice Description (editable)", value=PRESET_CHARACTERS["Male American"]["description"], lines=2)
-            emotion_select = gr.Dropdown(label="Select Emotion", choices=EMOTION_TAGS, value="<neutral>")
-            gen_btn = gr.Button("🔊 Generate Audio (LoRA)")
-            gen_logs = gr.Textbox(label="Logs", lines=10)
-        with gr.Column(scale=2):
-            gr.Markdown("### Output")
-            audio_player = gr.Audio(label="Generated Audio", type="filepath")
-            download_file = gr.File(label="Download generated file")
-            gr.Markdown("### Example")
-            gr.Textbox(label="Example Text", value=DEFAULT_TEXT, lines=2, interactive=False)
-            # example audio in repo: audio.wav (user can add)
-            example_audio_path = "audio.wav"
-            gr.Audio(label="Example Audio (project)", value=example_audio_path, type="filepath", interactive=False)
-
-    # wire updates: preset -> description
-    def _update_desc(preset_name):
-        return PRESET_CHARACTERS.get(preset_name, {}).get("description", "")
-    preset_select.change(fn=_update_desc, inputs=[preset_select], outputs=[description_box])
-
-    # generation
-    def _generate(text, preset, description, emotion):
-        return generate_for_ui(text, preset, description, emotion)
-
-    gen_btn.click(fn=_generate,
-                  inputs=[text_in, preset_select, description_box, emotion_select],
-                  outputs=[audio_player, download_file, gen_logs])
-
-    gr.Markdown("### Notes\n- Model & LoRA are loaded once at startup.\n- If you want faster generation on GPU, enable a GPU with `bitsandbytes` installed for 4-bit mode.\n- You can edit description to steer voice (sex/age/timbre).")
-
-if __name__ == "__main__":
-    demo.launch()
+
+# =========================================================
+# GRADIO UI
+# =========================================================
+with gr.Blocks() as demo:
+    gr.Markdown("## ⚡ Super-Fast Maya TTS (LoRA Merged + Optimized)")
+
+    txt = gr.Textbox(label="Enter text", value=DEFAULT_TEXT, lines=2)
+    btn = gr.Button("Generate Audio ⚡")
+    audio = gr.Audio(label="Audio Output", type="filepath")
+    file = gr.File(label="Download")
+    logs = gr.Textbox(label="Logs", lines=8)
+
+    btn.click(generate_audio, [txt], [audio, file, logs])
+
+    # Example section
+    gr.Markdown("### Example")
+    gr.Textbox(value=DEFAULT_TEXT, label="Example Text", interactive=False)
+    gr.Audio(value="audio.wav", type="filepath", label="Example Audio", interactive=False)
+
+demo.launch()