Update app.py

app.py

@@ -1,48 +1,53 @@
 import gradio as gr
 import torch
 from transformers import AutoProcessor, AutoModelForImageTextToText
-import nest_asyncio
-nest_asyncio.apply()
 
 # ---------------- MODEL SETUP ---------------- #
 MODEL_ID = "EpistemeAI/Audiogemma-3N-finetune"
 
+print("Loading processor...")
 processor = AutoProcessor.from_pretrained(MODEL_ID)
+
+print("Loading model...")
 model = AutoModelForImageTextToText.from_pretrained(
     MODEL_ID,
-    torch_dtype=torch.bfloat16,
+    torch_dtype=torch.float16,   # safer than bfloat16 on most GPUs
     device_map="auto"
 )
 
-# ---------------- TRANSLATION FUNCTION ---------------- #
+model.eval()
 
-def transcribe_and_translate(audio_input, target_language):
+# ---------------- TRANSLATION FUNCTION ---------------- #
 
+def transcribe_and_translate(audio_path, target_language):
     messages = [
         {
             "role": "user",
             "content": [
-                {"type": "audio", "audio": audio_input},
+                {"type": "audio", "audio": audio_path},
                 {
                     "type": "text",
-                    "text": f"Transcribe this audio into English, and then translate it into {target_language}."
+                    "text": f"Transcribe this audio into English, then translate it into {target_language}."
                 },
             ]
         }
     ]
 
-    input_ids = processor.apply_chat_template(
+    inputs = processor.apply_chat_template(
         messages,
         add_generation_prompt=True,
         tokenize=True,
-        return_dict=True,
-        return_tensors="pt",
+        return_tensors="pt"
     )
 
-    input_ids = input_ids.to(model.device, dtype=model.dtype)
+    inputs = inputs.to(model.device)
 
     with torch.no_grad():
-        outputs = model.generate(**input_ids, max_new_tokens=256)
+        outputs = model.generate(
+            **inputs,
+            max_new_tokens=256,
+            do_sample=False
+        )
 
     text = processor.batch_decode(
         outputs,
@@ -63,23 +68,14 @@ LANGUAGES = [
 
 with gr.Blocks() as demo:
     gr.Markdown("## 🎙️ Multilingual Audio Translator")
-    gr.Markdown("Speak English. The model will transcribe and translate into your chosen language.")
+    gr.Markdown("Upload or record English audio. The model will transcribe and translate it.")
 
     with gr.Row():
         audio_input = gr.Audio(type="filepath", label="Upload or Record Audio")
-        language_dropdown = gr.Dropdown(
-            choices=LANGUAGES,
-            value="French",
-            label="Target Language"
-        )
+        language_dropdown = gr.Dropdown(choices=LANGUAGES, value="French", label="Target Language")
 
     translate_btn = gr.Button("Translate")
-
-    output_text = gr.Textbox(
-        label="Translation Output",
-        lines=10,
-        interactive=False
-    )
+    output_text = gr.Textbox(label="Translation Output", lines=10)
 
     translate_btn.click(
         fn=transcribe_and_translate,
@@ -87,4 +83,9 @@ with gr.Blocks() as demo:
         outputs=output_text
     )
 
-demo.launch(debug=True)
+# ---------------- LAUNCH ---------------- #
+
+demo.launch(
+    server_port=7861,   # avoid stuck 7860
+    debug=True
+)