Update app_quant.py

app_quant.py

@@ -1,186 +1,192 @@
-# ---------------------------------------------------------
-#   Nava Ultra-Fast CPU Inference (4-bit Quant + Caching)
-# ---------------------------------------------------------
+# app_quant_fixed.py
 import gradio as gr
 import torch
 import soundfile as sf
 from pathlib import Path
+import traceback
+import time
 
-from transformers import (
-    AutoTokenizer,
-    AutoModelForCausalLM,
-    BitsAndBytesConfig
-)
-from peft import PeftModel
+from transformers import AutoTokenizer, AutoModelForCausalLM
+from peft import PeftModel, LoraConfig
 from snac import SNAC
 
-# ---------------------------------------------------------
-# CONFIG
-# ---------------------------------------------------------
+# -------------------------
 MODEL_NAME = "rahul7star/nava1.0"
 LORA_NAME = "rahul7star/nava-audio"
 SNAC_MODEL_NAME = "rahul7star/nava-snac"
 
-
-
 TARGET_SR = 24000
-DEFAULT_BATCH_SIZE = 500
-MICRO_BATCH = 2
-SEQ_LEN = 2048
 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…")
+DEFAULT_TEXT = "राजनीतिज्ञों ने कहा कि उन्होंने निर्णायक मत को अनावश्यक रूप से निर्धारित करने के लिए अफ़गान संविधान में काफी अस्पष्टता पाई थी"
+
+# conservative defaults
+SEQ_LEN_GPU = 240000   # if you really have GPU
+SEQ_LEN_CPU = 4096     # keep CPU small to avoid OOM
+MAX_NEW_TOKENS_CPU = 1024
+MAX_NEW_TOKENS_GPU = 240000
+
+# detect device
+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:
+        from transformers import BitsAndBytesConfig
+        bnb_available = True
+    else:
+        bnb_available = False
+except Exception:
+    bnb_available = False
+
+print(f"[init] CUDA available: {HAS_CUDA}, bitsandbytes available: {bnb_available}")
+
+# -------------------------
+# Load tokenizer (always)
+# -------------------------
+print("[init] Loading tokenizer...")
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True)
+
+# -------------------------
+# 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.")
+
+model.eval()
+
+# -------------------------
+# Load SNAC (once)
+# -------------------------
+print("[init] Loading SNAC...")
 snac_model = SNAC.from_pretrained(SNAC_MODEL_NAME).eval().to(DEVICE)
+print("[init] SNAC loaded.")
 
-
-# =========================================================
-#              INFERENCE FUNCTION
-# =========================================================
-def generate_audio_cpu_lora(text):
-
+# -------------------------
+# Inference function
+# -------------------------
+def generate_audio_cpu_lora(text: str):
     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
-# =========================================================
+    t0 = time.time()
+    try:
+        logs.append(f"[INFO] Device: {DEVICE} | SEQ_LEN: {SEQ_LEN} | MAX_NEW_TOKENS: {MAX_NEW_TOKENS}")
+
+        # 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
+
+        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
+
+        with torch.inference_mode():
+            outputs = model.generate(
+                **inputs,
+                max_new_tokens=max_new,
+                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
+            )
+
+        # extract generated part
+        gen_ids = outputs[0, inputs['input_ids'].shape[1]:].tolist()
+        logs.append(f"[INFO] Generated {len(gen_ids)} tokens")
+
+        # filter SNAC tokens (same logic)
+        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]
+
+        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.")
+            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)]
+
+        with torch.inference_mode():
+            z_q = snac_model.quantizer.from_codes(codes_tensor)
+            audio = snac_model.decoder(z_q)[0,0].cpu().numpy()
+
+        if len(audio) > 2048:
+            audio = audio[2048:]
+
+        out_path = OUT_ROOT / f"tts_output_cpu_lora.wav"
+        sf.write(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")
+        return str(out_path), str(out_path), "\n".join(logs)
+
+    except Exception as e:
+        tb = traceback.format_exc()
+        logs.append(f"[ERROR] {e}\n{tb}")
+        return None, None, "\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)
+    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 = gr.Textbox(label="Logs")
-
-    btn.click(generate_audio_cpu_lora, [txt], [audio, file, logs])
-
+    logs_box = gr.Textbox(label="Logs", lines=10)
+    btn.click(generate_audio_cpu_lora, [txt], [audio, file, logs_box])
 
 if __name__ == "__main__":
     demo.launch()