Update app.py

app.py

@@ -4,7 +4,7 @@ import gradio as gr
 import torch
 import torchaudio
 import spaces
-import nemo.collections.asr as nemo_asr
+from transformers import AutoProcessor, AutoModelForSpeechSeq2Seq, AutoModelForCTC
 
 LANGUAGE_NAME_TO_CODE = {
     "Assamese": "as", "Bengali": "bn", "Bodo": "br", "Dogri": "doi",
@@ -15,30 +15,43 @@ LANGUAGE_NAME_TO_CODE = {
     "Telugu": "te", "Urdu": "ur"
 }
 
-DESCRIPTION = """IndicConformer: Dual-Decoder ASR for Indian Languages"""
+DESCRIPTION = "IndicConformer-600M Multilingual ASR (CTC + RNNT)"
 
-device = "cuda:0" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
-model = nemo_asr.models.EncDecCTCModel.from_pretrained("ai4bharat/IndicConformer").to(device)
-model.eval()
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+# Load processor and models
+processor = AutoProcessor.from_pretrained("ai4bharat/indic-conformer-600m-multilingual", trust_remote_code=True)
+
+model_ctc = AutoModelForCTC.from_pretrained("ai4bharat/indic-conformer-600m-multilingual", trust_remote_code=True).to(device)
+model_ctc.eval()
+
+model_rnnt = AutoModelForSpeechSeq2Seq.from_pretrained("ai4bharat/indic-conformer-600m-multilingual", trust_remote_code=True).to(device)
+model_rnnt.eval()
 
 @spaces.GPU
 def transcribe_ctc_and_rnnt(audio_path, language_name):
     lang_id = LANGUAGE_NAME_TO_CODE[language_name]
-    waveform, sample_rate = torchaudio.load(audio_path)
+
+    waveform, sr = torchaudio.load(audio_path)
     waveform = waveform.mean(dim=0, keepdim=True) if waveform.shape[0] > 1 else waveform
-    waveform = torchaudio.functional.resample(waveform, sample_rate, 16000)
-    waveform_np = waveform.squeeze().numpy()
+    waveform = torchaudio.functional.resample(waveform, sr, 16000)
+
+    input_values = processor(waveform.squeeze().numpy(), sampling_rate=16000, return_tensors="pt").input_values.to(device)
 
-    model.cur_decoder = "ctc"
-    ctc = model.transcribe([waveform_np], batch_size=1, language_id=lang_id)[0][0]
+    with torch.no_grad():
+        # CTC decoding
+        ctc_logits = model_ctc(input_values).logits
+        ctc_ids = torch.argmax(ctc_logits, dim=-1)
+        ctc_output = processor.batch_decode(ctc_ids)[0]
 
-    model.cur_decoder = "rnnt"
-    rnnt = model.transcribe([waveform_np], batch_size=1, language_id=lang_id)[0][0]
+        # RNNT decoding
+        rnnt_output = processor.batch_decode(model_rnnt.generate(input_values, decoder_input_ids=torch.tensor([[processor.tokenizer.lang2id[lang_id]]]).to(device)))[0]
 
-    return ctc, rnnt
+    return ctc_output.strip(), rnnt_output.strip()
 
+# Gradio interface
 with gr.Blocks() as demo:
-    gr.Markdown(DESCRIPTION)
+    gr.Markdown(f"## {DESCRIPTION}")
     with gr.Row():
         with gr.Column():
             audio = gr.Audio(label="Upload or record audio", type="filepath")