Update app.py

app.py

@@ -1,10 +1,11 @@
 from __future__ import annotations
-import os
-import gradio as gr
 import torch
 import torchaudio
+import gradio as gr
 import spaces
-from transformers import AutoProcessor, AutoModelForSpeechSeq2Seq, AutoModelForCTC
+from transformers import AutoModel
+
+DESCRIPTION = "IndicConformer-600M Multilingual ASR (CTC + RNNT)"
 
 LANGUAGE_NAME_TO_CODE = {
     "Assamese": "as", "Bengali": "bn", "Bodo": "br", "Dogri": "doi",
@@ -15,55 +16,48 @@ LANGUAGE_NAME_TO_CODE = {
     "Telugu": "te", "Urdu": "ur"
 }
 
-DESCRIPTION = "IndicConformer-600M Multilingual ASR (CTC + RNNT)"
-
 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()
+# Load Indic Conformer model (assumes custom forward handles decoding strategy)
+model = AutoModel.from_pretrained("ai4bharat/indic-conformer-600m-multilingual", trust_remote_code=True).to(device)
+model.eval()
 
 @spaces.GPU
 def transcribe_ctc_and_rnnt(audio_path, language_name):
-    lang_id = LANGUAGE_NAME_TO_CODE[language_name]
+    lang_code = LANGUAGE_NAME_TO_CODE[language_name]
 
+    # Load and preprocess audio
     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, sr, 16000)
-
-    input_values = processor(waveform.squeeze().numpy(), sampling_rate=16000, return_tensors="pt").input_values.to(device)
-
-    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]
+    waveform = torchaudio.functional.resample(waveform, sr, 16000).to(device)
 
-        # 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]
+    try:
+        # Assume model's forward method takes waveform, language code, and decoding type
+        with torch.no_grad():
+            transcription_ctc = model(waveform, lang_code, "ctc")
+            transcription_rnnt = model(waveform, lang_code, "rnnt")
+    except Exception as e:
+        return f"Error: {str(e)}", ""
 
-    return ctc_output.strip(), rnnt_output.strip()
+    return transcription_ctc.strip(), transcription_rnnt.strip()
 
-# Gradio interface
+# Gradio UI
 with gr.Blocks() as demo:
     gr.Markdown(f"## {DESCRIPTION}")
     with gr.Row():
         with gr.Column():
-            audio = gr.Audio(label="Upload or record audio", type="filepath")
+            audio = gr.Audio(label="Upload or Record Audio", type="filepath")
             lang = gr.Dropdown(
-                label="Select language",
-                choices=LANGUAGE_NAME_TO_CODE.keys(),
+                label="Select Language",
+                choices=list(LANGUAGE_NAME_TO_CODE.keys()),
                 value="Hindi"
             )
             transcribe_btn = gr.Button("Transcribe (CTC + RNNT)")
         with gr.Column():
-            ctc_output = gr.Textbox(label="CTC Transcription")
-            rnnt_output = gr.Textbox(label="RNNT Transcription")
+            gr.Markdown("### CTC Transcription")
+            ctc_output = gr.Textbox(lines=3)
+            gr.Markdown("### RNNT Transcription")
+            rnnt_output = gr.Textbox(lines=3)
 
     transcribe_btn.click(fn=transcribe_ctc_and_rnnt, inputs=[audio, lang], outputs=[ctc_output, rnnt_output])