Update app.py

app.py

@@ -1,32 +1,95 @@
-import gradio as gr
+import os
 import torch
 import spaces
-from transformers import AutoModelForCausalLM, AutoTokenizer
-import scipy.io.wavfile as wavfile
-
-# Model loads on HF Server, 0% of your local data used
-model_id = "maya-research/Veena"
-tokenizer = AutoTokenizer.from_pretrained(model_id)
-model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.bfloat16)
-
-@spaces.GPU  # Essential for when you get the GPU grant
-def generate_audio(text):
-    inputs = tokenizer(text, return_tensors="pt")
-    # Veena generation logic
+import gradio as gr
+import soundfile as sf
+from snac import SNAC
+from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
+
+# --- CONFIGURATION ---
+MODEL_ID = "maya-research/Veena"
+SNAC_MODEL_ID = "hubertsiuzdak/snac_24khz"
+VALID_KEY = os.environ.get("MY_API_KEY") # Set this in HF Space Secrets
+
+# Token Offsets for Veena
+START_OF_SPEECH_TOKEN = 128257
+END_OF_SPEECH_TOKEN = 128258
+START_OF_HUMAN_TOKEN = 128259
+END_OF_HUMAN_TOKEN = 128260
+START_OF_AI_TOKEN = 128261
+END_OF_AI_TOKEN = 128262
+AUDIO_CODE_BASE_OFFSET = 128266
+
+# --- MODEL LOADING ---
+# 4-bit config allows it to run on smaller/shared GPUs
+quant_config = BitsAndBytesConfig(
+    load_in_4bit=True,
+    bnb_4bit_quant_type="nf4",
+    bnb_4bit_compute_dtype=torch.bfloat16,
+)
+
+tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL_ID, 
+    quantization_config=quant_config, 
+    device_map="auto"
+)
+snac_model = SNAC.from_pretrained(SNAC_MODEL_ID).eval().to("cuda" if torch.cuda.is_available() else "cpu")
+
+def decode_audio(tokens):
+    """Converts Veena's tokens into a WAV file"""
+    snac_tokens = [t for t in tokens if t >= AUDIO_CODE_BASE_OFFSET]
+    if not snac_tokens or len(snac_tokens) % 7 != 0:
+        return None
+        
+    codes_lvl = [[] for _ in range(3)]
+    # De-interleave based on Veena's 7-token frame structure
+    for i in range(0, len(snac_tokens), 7):
+        codes_lvl[0].append(snac_tokens[i] - AUDIO_CODE_BASE_OFFSET)
+        codes_lvl[1].extend([snac_tokens[i+1]- (AUDIO_CODE_BASE_OFFSET + 4096), snac_tokens[i+2]- (AUDIO_CODE_BASE_OFFSET + 8192)])
+        codes_lvl[2].extend([snac_tokens[i+3]- (AUDIO_CODE_BASE_OFFSET + 12288), snac_tokens[i+4]- (AUDIO_CODE_BASE_OFFSET + 16384), 
+                             snac_tokens[i+5]- (AUDIO_CODE_BASE_OFFSET + 20480), snac_tokens[i+6]- (AUDIO_CODE_BASE_OFFSET + 24576)])
+    
+    codes = [torch.tensor([c]).to(snac_model.device) for c in codes_lvl]
+    with torch.no_grad():
+        audio_values = snac_model.decode(codes)
+    return audio_values.cpu().numpy().squeeze()
+
+@spaces.GPU
+def generate_veena_speech(text, api_key, speaker="kavya"):
+    # Security check for n8n
+    if api_key != VALID_KEY:
+        raise gr.Error("Invalid API Key")
+
+    # Format prompt for Veena
+    prompt = [START_OF_HUMAN_TOKEN] + tokenizer.encode(f"<spk_{speaker}> {text}") + [END_OF_HUMAN_TOKEN, START_OF_AI_TOKEN]
+    input_ids = torch.tensor([prompt]).to(model.device)
+    
     with torch.no_grad():
-        output_tokens = model.generate(**inputs, max_new_tokens=512)
+        output = model.generate(
+            input_ids, 
+            max_new_tokens=1024, 
+            do_sample=True, 
+            eos_token_id=[END_OF_SPEECH_TOKEN, END_OF_AI_TOKEN]
+        )
     
-    # Save to a temporary file for the API to return
-    filename = "output.wav"
-    # (Assuming Veena's standard output processing here)
-    # wavfile.write(filename, 24000, audio_data) 
-    return filename
+    audio_data = decode_audio(output[0].tolist())
+    if audio_data is not None:
+        output_path = "output.wav"
+        sf.write(output_path, audio_data, 24000)
+        return output_path
+    return None
 
+# --- GRADIO INTERFACE ---
 demo = gr.Interface(
-    fn=generate_audio,
-    inputs=gr.Textbox(label="Input Text"),
-    outputs=gr.Audio(type="filepath"),
-    api_name="predict" # This creates the /predict endpoint for n8n
+    fn=generate_veena_speech,
+    inputs=[
+        gr.Textbox(label="Text to Speak"),
+        gr.Textbox(label="API Key", type="password"),
+        gr.Dropdown(choices=["kavya", "agastya", "maitri", "vinaya"], value="kavya", label="Speaker")
+    ],
+    outputs=gr.Audio(label="Generated Audio"),
+    api_name="predict"
 )
 
 demo.launch()