Update app_quant.py

app_quant.py

@@ -1,122 +1,264 @@
-# app_quant_fixed.py
+# 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
+
 import gradio as gr
-import torch
+import numpy as np
 import soundfile as sf
-from pathlib import Path
-import traceback
-import time
+import torch
 
 from transformers import AutoTokenizer, AutoModelForCausalLM
 from peft import PeftModel, LoraConfig
 from snac import SNAC
 
+# -------------------------
+# Config - update if needed
 # -------------------------
 MODEL_NAME = "rahul7star/nava1.0"
 LORA_NAME = "rahul7star/nava-audio"
-SNAC_MODEL_NAME = "rahul7star/nava-snac"
-
+SNAC_MODEL_NAME = "hubertsiuzdak/snac_24khz"  # or your snac model name
 TARGET_SR = 24000
 OUT_ROOT = Path("/tmp/data")
 OUT_ROOT.mkdir(exist_ok=True, parents=True)
 
-DEFAULT_TEXT = "राजनीतिज्ञों ने कहा कि उन्होंने निर्णायक मत को अनावश्यक रूप से निर्धारित करने के लिए अफ़गान संविधान में काफी अस्पष्टता पाई थी"
-
-# conservative defaults
-SEQ_LEN_GPU = 240000   # if you really have GPU
-SEQ_LEN_CPU = 4096     # keep CPU small to avoid OOM
+# 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
+# -------------------------
+# Detect device & bitsandbytes
+# -------------------------
 HAS_CUDA = torch.cuda.is_available()
 DEVICE = "cuda" if HAS_CUDA else "cpu"
 
-# optional: try import bitsandbytes only if CUDA available
-try:
-    if HAS_CUDA:
+bnb_available = False
+if HAS_CUDA:
+    try:
         from transformers import BitsAndBytesConfig
         bnb_available = True
-    else:
+    except Exception:
         bnb_available = False
-except Exception:
-    bnb_available = False
 
-print(f"[init] CUDA available: {HAS_CUDA}, bitsandbytes available: {bnb_available}")
+print(f"[init] Device: {DEVICE} | CUDA: {HAS_CUDA} | bnb: {bnb_available}")
 
 # -------------------------
-# Load tokenizer (always)
+# Load tokenizer
 # -------------------------
 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 (GPU vs CPU safe)
-# -------------------------
-print("[init] Loading base model + LoRA (this may take a while)...")
-if HAS_CUDA and bnb_available:
-    # GPU + bnb path: use 4-bit quant
-    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,
-    )
-    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] Loaded model in 4-bit (GPU).")
-else:
-    # CPU fallback: load in FP32 with low_cpu_mem_usage
-    # Avoid load_in_4bit on CPU
-    base_model = AutoModelForCausalLM.from_pretrained(
-        MODEL_NAME,
-        torch_dtype=torch.float32,
-        device_map={"": "cpu"},
-        low_cpu_mem_usage=True,
-        trust_remote_code=True,
-    )
-    # attach PEFT adapter - this will add LoRA wrappers but keep base weights on CPU
-    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] Loaded model on CPU (FP32) with LoRA.")
+# 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
 
 model.eval()
 
 # -------------------------
-# Load SNAC (once)
+# Load SNAC Model (once)
 # -------------------------
-print("[init] Loading SNAC...")
+print("[init] Loading SNAC decoder...")
 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.")
 
 # -------------------------
-# Inference function
+# 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_cpu_lora(text: str):
+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
+    """
     logs = []
     t0 = time.time()
     try:
-        logs.append(f"[INFO] Device: {DEVICE} | SEQ_LEN: {SEQ_LEN} | MAX_NEW_TOKENS: {MAX_NEW_TOKENS}")
+        # 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
 
-        # Build prompt (same as your earlier code)
-        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
+        logs.append(f"[INFO] DEVICE={DEVICE} | prompt_len={len(prompt)}")
 
+        # Tokenize
         inputs = tokenizer(prompt, return_tensors="pt", padding=True).to(DEVICE)
 
-        # Keep generated tokens small on CPU
-        max_new = min(MAX_NEW_TOKENS, 1024) if DEVICE == "cpu" else MAX_NEW_TOKENS
+        # 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)
 
+        # Generate
         with torch.inference_mode():
             outputs = model.generate(
                 **inputs,
@@ -126,49 +268,45 @@ def generate_audio_cpu_lora(text: str):
                 repetition_penalty=1.1,
                 do_sample=True,
                 eos_token_id=128258,
-                pad_token_id=tokenizer.pad_token_id
+                pad_token_id=tokenizer.pad_token_id,
             )
 
-        # extract generated part
-        gen_ids = outputs[0, inputs['input_ids'].shape[1]:].tolist()
+        gen_ids = outputs[0, inputs["input_ids"].shape[1]:].tolist()
         logs.append(f"[INFO] Generated {len(gen_ids)} tokens")
 
-        # filter SNAC tokens (same logic)
+        # 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]
 
+        # make frames multiple of 7
         frames = len(snac_tokens) // 7
-        snac_tokens = snac_tokens[:frames*7]
-
-        l1, l2, l3 = [], [], []
-        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])
-
-        if len(l1) == 0:
-            logs.append("[WARN] No SNAC frames found in generated tokens. Returning debug logs.")
+        if frames == 0:
+            logs.append("[WARN] No SNAC frames found in generated tokens.")
             return None, None, "\n".join(logs)
 
-        codes_tensor = [torch.tensor(l1, dtype=torch.long, device=DEVICE).unsqueeze(0),
-                        torch.tensor(l2, dtype=torch.long, device=DEVICE).unsqueeze(0),
-                        torch.tensor(l3, dtype=torch.long, device=DEVICE).unsqueeze(0)]
+        l1, l2, l3 = snac_tokens_to_levels(snac_tokens)
+        if l1 is None:
+            logs.append("[ERROR] SNAC conversion failed.")
+            return None, None, "\n".join(logs)
 
+        # decode via SNAC (keep on DEVICE)
         with torch.inference_mode():
-            z_q = snac_model.quantizer.from_codes(codes_tensor)
-            audio = snac_model.decoder(z_q)[0,0].cpu().numpy()
+            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()
 
+        # trim warmup samples
         if len(audio) > 2048:
             audio = audio[2048:]
 
-        out_path = OUT_ROOT / f"tts_output_cpu_lora.wav"
-        sf.write(out_path, 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"[TIME] Elapsed {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)
 
     except Exception as e:
@@ -176,28 +314,63 @@ def generate_audio_cpu_lora(text: str):
         logs.append(f"[ERROR] {e}\n{tb}")
         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() as demo:
-    gr.Markdown("## Maya TTS — CPU/GPU safe")
-    txt = gr.Textbox(label="Enter text", value=DEFAULT_TEXT, lines=2)
-    btn = gr.Button("Generate Audio")
-    audio = gr.Audio(label="Audio", type="filepath")
-    file = gr.File(label="Download")
-    logs_box = gr.Textbox(label="Logs", lines=10)
-    btn.click(generate_audio_cpu_lora, [txt], [audio, file, logs_box])
-
-
-      # -----------------------------
-    # Example section
-    # -----------------------------
-    gr.Markdown("### Example")
-    example_text = DEFAULT_TEXT
-    example_audio_path = "audio.wav" 
-
-    gr.Textbox(label="Example Text", value=example_text, lines=2, interactive=False)
-    gr.Audio(label="Example Audio", value=example_audio_path, type="filepath", interactive=False)
+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()