Update app.py

app.py

@@ -1,27 +1,16 @@
-import os
 import torch
-import torchaudio
-import gradio as gr
 from transformers import AutoProcessor, AutoModelForImageTextToText
 
-# Load model
+# Load processor and model
 processor = AutoProcessor.from_pretrained("EpistemeAI/Audiogemma-3N-finetune")
 model = AutoModelForImageTextToText.from_pretrained(
     "EpistemeAI/Audiogemma-3N-finetune",
-    torch_dtype=torch.bfloat16,
-).to("cuda")
-
-
-
+    torch_dtype=torch.bfloat16,  # Match model weights to bfloat16
+    device_map="auto"             # Sends model to GPU if available
+)
 
 def convert_audio_to_text(audio_file):
-    # Load audio
-    waveform, sample_rate = torchaudio.load(audio_file)
-
-    # Convert to mono if stereo
-    if waveform.shape[0] > 1:
-        waveform = waveform.mean(dim=0, keepdim=True)
-
+    # Build the multimodal chat message
     messages = [
         {
             "role": "user",
@@ -32,36 +21,34 @@ def convert_audio_to_text(audio_file):
         }
     ]
 
-    input_ids = processor.apply_chat_template(
+    # Prepare inputs
+    inputs = processor.apply_chat_template(
         messages,
         add_generation_prompt=True,
-        tokenize=True, return_dict=True,
+        tokenize=True,
+        return_dict=True,
         return_tensors="pt",
     ).to("cuda")
 
-    outputs = model.generate(**input_ids, max_new_tokens=512)
+    # Cast all floating tensors to match model dtype
+    for k, v in inputs.items():
+        if torch.is_floating_point(v):
+            inputs[k] = v.to(dtype=torch.bfloat16)
+
+    # Generate output tokens
+    outputs = model.generate(**inputs, max_new_tokens=512)
 
+    # Decode output
     result = processor.batch_decode(
         outputs,
-        skip_special_tokens=False,
-        clean_up_tokenization_spaces=False
+        skip_special_tokens=True,
+        clean_up_tokenization_spaces=True
     )
-    return result
-
-
-with gr.Blocks() as demo:
-    gr.Markdown("## 🎙️ Audio Transcription with Audiogemma")
 
-    with gr.Tab("Microphone"):
-        mic = gr.Audio(sources="microphone", type="filepath")
-        mic_out = gr.Textbox()
-        mic_btn = gr.Button("Transcribe")
-        mic_btn.click(convert_audio_to_text, mic, mic_out)
-
-    with gr.Tab("Audio File"):
-        file = gr.Audio(sources="upload", type="filepath")
-        file_out = gr.Textbox()
-        file_btn = gr.Button("Transcribe")
-        file_btn.click(convert_audio_to_text, file, file_out)
+    return result
 
-demo.launch()
+# Example usage
+if __name__ == "__main__":
+    audio_path = "sample_audio.wav"
+    transcription_translation = convert_audio_to_text(audio_path)
+    print(transcription_translation)