Update app.py

app.py

@@ -1,105 +1,90 @@
-import os
-os.environ["TORCHDYNAMO_DISABLE"] = "1"
-
 import gradio as gr
 import torch
-import librosa
-import numpy as np
-from threading import Thread
-from unsloth import FastModel
-from transformers import AutoProcessor, TextIteratorStreamer
-
-TARGET_SAMPLING_RATE = 16000
-device = "cuda" if torch.cuda.is_available() else "cpu"
-
-print("Loading Gemma-3N audio model...")
-
-# IMPORTANT: disable alt-up (fixes uint8 clamp crash)
-model, tokenizer = FastModel.from_pretrained(
-    model_name="unsloth/gemma-3n-E4B-it-unsloth-bnb-4bit",
-    max_seq_length=2048,
-    dtype=None,
-    load_in_4bit=True,
-    full_finetuning=False,
-    device_map="auto",
-)
-
-processor = AutoProcessor.from_pretrained(
-    "EpistemeAI/Audiogemma-3N-finetune"
+from transformers import AutoProcessor, AutoModelForImageTextToText
+import nest_asyncio
+nest_asyncio.apply()
+
+# ---------------- MODEL SETUP ---------------- #
+MODEL_ID = "EpistemeAI/Audiogemma-3N-finetune"
+
+processor = AutoProcessor.from_pretrained(MODEL_ID)
+model = AutoModelForImageTextToText.from_pretrained(
+    MODEL_ID,
+    torch_dtype=torch.bfloat16,
+    device_map="auto"
 )
 
-model.eval()
-print("Model loaded on", device)
+# ---------------- TRANSLATION FUNCTION ---------------- #
 
-# ---------------- AUDIO PIPELINE ---------------- #
-
-def transcribe_and_translate(audio_input):
-    if audio_input is None:
-        yield "Please upload or record audio."
-        return
+def transcribe_and_translate(audio_input, target_language):
 
     messages = [
-        {
-            "role": "system",
-            "content": [
-                {"type": "text", "text": "You transcribe spoken audio and translate it into German."}
-            ],
-        },
         {
             "role": "user",
             "content": [
                 {"type": "audio", "audio": audio_input},
-                {"type": "text", "text": "Please transcribe this audio and translate it to German."}
-            ],
-        },
+                {
+                    "type": "text",
+                    "text": f"Transcribe this audio into English, and then translate it into {target_language}."
+                },
+            ]
+        }
     ]
 
-    inputs = processor.apply_chat_template(
+    input_ids = processor.apply_chat_template(
         messages,
-        tokenize=True,
         add_generation_prompt=True,
-        return_tensors="pt",
+        tokenize=True,
         return_dict=True,
-    ).to(device)
+        return_tensors="pt",
+    )
 
-    streamer = TextIteratorStreamer(processor, skip_prompt=True)
+    input_ids = input_ids.to(model.device, dtype=model.dtype)
 
-    generation_kwargs = dict(
-        **inputs,
-        max_new_tokens=1024,
-        temperature=0.7,
-        top_p=0.95,
-        top_k=50,
-        streamer=streamer,
-    )
+    with torch.no_grad():
+        outputs = model.generate(**input_ids, max_new_tokens=256)
 
-    thread = Thread(target=model.generate, kwargs=generation_kwargs)
-    thread.start()
+    text = processor.batch_decode(
+        outputs,
+        skip_special_tokens=True,
+        clean_up_tokenization_spaces=True
+    )
 
-    output = ""
-    for token in streamer:
-        output += token
-        yield output
+    return text[0]
 
 
 # ---------------- GRADIO UI ---------------- #
 
-with gr.Blocks(theme=gr.themes.Soft()) as demo:
-    gr.Markdown("# Gemma-3N Audio Transcription + German Translation")
+LANGUAGES = [
+    "French", "Spanish", "German", "Italian", "Portuguese",
+    "Chinese", "Japanese", "Korean", "Arabic", "Hindi",
+    "Russian", "Ukrainian", "Hebrew", "Thai", "Vietnamese"
+]
+
+with gr.Blocks() as demo:
+    gr.Markdown("## 🎙️ Multilingual Audio Translator")
+    gr.Markdown("Speak English. The model will transcribe and translate into your chosen language.")
 
     with gr.Row():
-        audio_input = gr.Audio(
-            sources=["upload", "microphone"],
-            type="numpy",
-            label="Audio Input"
-        )
-        text_output = gr.Textbox(
-            label="Transcription + Translation",
-            lines=12
+        audio_input = gr.Audio(type="filepath", label="Upload or Record Audio")
+        language_dropdown = gr.Dropdown(
+            choices=LANGUAGES,
+            value="French",
+            label="Target Language"
         )
 
-    btn = gr.Button("Transcribe and Translate", variant="primary")
-    btn.click(transcribe_and_translate, audio_input, text_output)
+    translate_btn = gr.Button("Translate")
+
+    output_text = gr.Textbox(
+        label="Translation Output",
+        lines=10,
+        interactive=False
+    )
+
+    translate_btn.click(
+        fn=transcribe_and_translate,
+        inputs=[audio_input, language_dropdown],
+        outputs=output_text
+    )
 
-if __name__ == "__main__":
-    demo.launch()
+demo.launch(debug=True)