Update app.py

app.py

@@ -1,38 +1,43 @@
 import gradio as gr
-from transformers import pipeline, AutoModelForSeq2SeqLM, AutoTokenizer, AutoModelForCTC, Wav2Vec2Processor
+from transformers import AutoModel, AutoTokenizer, AutoModelForSeq2SeqLM, AutoModelForCTC
 import torch
+import torchaudio
 import librosa
 import soundfile as sf
 import io
+import os
 
-# Load models (use latest from AI4Bharat)
-asr_model_name = "ai4bharat/indicconformer-600m-multilingual"
-asr_processor = Wav2Vec2Processor.from_pretrained(asr_model_name)
-asr_model = AutoModelForCTC.from_pretrained(asr_model_name)
+# For gated models, set token
+os.environ["HF_TOKEN"] = "YOUR_HF_TOKEN"  # From huggingface.co/settings/tokens
 
-llm_model_name = "ai4bharat/IndicBART"  # Fine-tuned on IndicAlign
-llm_tokenizer = AutoTokenizer.from_pretrained(llm_model_name)
+# Load models
+asr_model_name = "ai4bharat/indic-conformer-600m-multilingual"
+asr_model = AutoModel.from_pretrained(asr_model_name, trust_remote_code=True)
+
+llm_model_name = "ai4bharat/IndicBART"
+llm_tokenizer = AutoTokenizer.from_pretrained(llm_model_name, do_lower_case=False, use_fast=False, keep_accents=True)
 llm_model = AutoModelForSeq2SeqLM.from_pretrained(llm_model_name)
 
 trans_model_name = "ai4bharat/IndicTrans3-beta"
 trans_tokenizer = AutoTokenizer.from_pretrained(trans_model_name)
 trans_model = AutoModelForSeq2SeqLM.from_pretrained(trans_model_name)
 
-tts_model_name = "ai4bharat/indicf5"  # Or indic-parler-tts-v2
-tts_pipe = pipeline("text-to-speech", model=tts_model_name)
+tts_model_name = "ai4bharat/IndicF5"
+tts_model = AutoModel.from_pretrained(tts_model_name, trust_remote_code=True)
 
 def full_pipeline(audio, source_lang, target_lang):
-    # ASR: Audio to text
-    audio_array, _ = librosa.load(io.BytesIO(audio), sr=16000)
-    inputs = asr_processor(audio_array, sampling_rate=16000, return_tensors="pt")
-    with torch.no_grad():
-        logits = asr_model(inputs.input_values).logits
-    pred_ids = torch.argmax(logits, dim=-1)
-    text = asr_processor.batch_decode(pred_ids)[0]
-
-    # LLM: Generate response (simple echo for now; enhance later)
-    inputs = llm_tokenizer(text, return_tensors="pt")
-    outputs = llm_model.generate(**inputs)
+    # ASR
+    audio_array, sr = librosa.load(io.BytesIO(audio), sr=16000)
+    wav = torch.tensor(audio_array).unsqueeze(0)
+    text = asr_model(wav, source_lang, "ctc")
+
+    # LLM: Simple generation
+    bos_id = llm_tokenizer._convert_token_to_id_with_added_voc("<s>")
+    eos_id = llm_tokenizer._convert_token_to_id_with_added_voc("</s>")
+    pad_id = llm_tokenizer._convert_token_to_id_with_added_voc("<pad>")
+    lang_code = f"<2{source_lang}>"  # e.g. <2hi>
+    inputs = llm_tokenizer(text + " </s> " + lang_code, add_special_tokens=False, return_tensors="pt")
+    outputs = llm_model.generate(**inputs, max_length=50, decoder_start_token_id=llm_tokenizer._convert_token_to_id_with_added_voc(lang_code))
     response = llm_tokenizer.decode(outputs[0], skip_special_tokens=True)
 
     # Translation if needed
@@ -41,17 +46,19 @@ def full_pipeline(audio, source_lang, target_lang):
         outputs = trans_model.generate(**inputs)
         response = trans_tokenizer.decode(outputs[0], skip_special_tokens=True)
 
-    # TTS: Text to audio
-    tts_output = tts_pipe(response)
+    # TTS (needs ref audio; use example)
+    ref_audio_path = "prompts/example.wav"  # Upload example prompt to repo
+    ref_text = "Example reference text in language"
+    tts_output = tts_model(response, ref_audio_path=ref_audio_path, ref_text=ref_text)
     with io.BytesIO() as buffer:
-        sf.write(buffer, tts_output["audio"][0], tts_output["sampling_rate"], format="wav")
+        sf.write(buffer, tts_output, 24000, format="wav")
         audio_bytes = buffer.getvalue()
 
     return audio_bytes, text, response
 
 iface = gr.Interface(
     fn=full_pipeline,
-    inputs=[gr.Audio(type="file"), gr.Textbox(label="Source Lang"), gr.Textbox(label="Target Lang")],
+    inputs=[gr.Audio(type="file"), gr.Textbox(label="Source Lang e.g. hi"), gr.Textbox(label="Target Lang e.g. en")],
     outputs=[gr.Audio(label="Response Audio"), gr.Textbox(label="Transcribed Text"), gr.Textbox(label="Response Text")],
     title="HanuVak Indic Conversation Backend"
 )