Update app.py

app.py

@@ -1,18 +1,18 @@
 import gradio as gr
-from transformers import AutoModel, AutoTokenizer, AutoModelForSeq2SeqLM, AutoModelForCTC
+from transformers import AutoProcessor, AutoModelForCTC, AutoModelForSeq2SeqLM, AutoTokenizer, pipeline
 import torch
-import torchaudio
 import librosa
 import soundfile as sf
 import io
 import os
 
-# For gated models, set token
-os.environ["HF_TOKEN"] = "YOUR_HF_TOKEN"  # From huggingface.co/settings/tokens
+# Use HF_TOKEN from env
+os.environ["HF_TOKEN"] = os.getenv("HF_TOKEN")
 
-# Load models
-asr_model_name = "ai4bharat/indic-conformer-600m-multilingual"
-asr_model = AutoModel.from_pretrained(asr_model_name, trust_remote_code=True)
+# Models (use CPU if no GPU; for free tier, may be slow/large - upgrade for GPU)
+asr_model_name = "ai4bharat/indicconformer-600m-multilingual"
+asr_processor = AutoProcessor.from_pretrained(asr_model_name)
+asr_model = AutoModelForCTC.from_pretrained(asr_model_name)
 
 llm_model_name = "ai4bharat/IndicBART"
 llm_tokenizer = AutoTokenizer.from_pretrained(llm_model_name, do_lower_case=False, use_fast=False, keep_accents=True)
@@ -22,36 +22,32 @@ 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"
-tts_model = AutoModel.from_pretrained(tts_model_name, trust_remote_code=True)
+tts_pipe = pipeline("text-to-speech", model="ai4bharat/indic-parler-tts-v2")  # Switch to non-gated if issues
 
 def full_pipeline(audio, source_lang, target_lang):
     # 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))
+    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 response (echo for test)
+    inputs = llm_tokenizer(text, return_tensors="pt")
+    outputs = llm_model.generate(**inputs)
     response = llm_tokenizer.decode(outputs[0], skip_special_tokens=True)
 
-    # Translation if needed
+    # Translation
     if source_lang != target_lang:
         inputs = trans_tokenizer(response, return_tensors="pt")
         outputs = trans_model.generate(**inputs)
         response = trans_tokenizer.decode(outputs[0], skip_special_tokens=True)
 
-    # 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)
+    # TTS
+    tts_output = tts_pipe(response)
     with io.BytesIO() as buffer:
-        sf.write(buffer, tts_output, 24000, format="wav")
+        sf.write(buffer, tts_output["audio"][0], tts_output["sampling_rate"], format="wav")
         audio_bytes = buffer.getvalue()
 
     return audio_bytes, text, response
@@ -60,7 +56,7 @@ iface = gr.Interface(
     fn=full_pipeline,
     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"
+    title="HanuVak Backend"
 )
 
 if __name__ == "__main__":