Add SeamlessExpressive + SeamlessM4T v2 Large + NLLB-200

README.md

@@ -11,14 +11,14 @@ license: mit
 hardware: t4-small
 ---
 
-# 🌍 STTR - Speech-to-Text & Translation API
+# 🌍 STTR - Speech & Translation API
 
-**Meta AI Models:**
+## Meta AI Models:
 - 🎤 **SeamlessM4T v2 Large** - STT (101 languages)
 - 🌍 **NLLB-200** - Translation (200 languages + Darija!)
-- 🎭 **SeamlessExpressive** - Expressive Speech Translation
+- 🎭 **SeamlessExpressive** - Expressive Speech Translation (preserves tone!)
 
-**API Endpoints:**
+## API Endpoints:
 - `/stt` - Speech-to-Text
 - `/translate` - Text Translation
-- `/expressive` - Expressive Speech Translation
+- `/expressive` - Expressive Speech-to-Speech Translation

app.py

@@ -3,10 +3,12 @@ from transformers import (
     AutoProcessor, 
     SeamlessM4Tv2ForSpeechToText,
     AutoModelForSeq2SeqLM, 
-    AutoTokenizer
+    AutoTokenizer,
+    SeamlessM4Tv2Model,
 )
 import torch
 import numpy as np
+import torchaudio
 
 # ============================================================
 # 🚀 Device Setup
@@ -19,24 +21,35 @@ print(f"🖥️ Device: {device}")
 # 📥 Load Models
 # ============================================================
 
-# SeamlessM4T v2 Large for STT
-print("📥 Loading SeamlessM4T v2 Large...")
+# 1. SeamlessM4T v2 Large for STT
+print("📥 Loading SeamlessM4T v2 Large (STT)...")
 STT_MODEL = "facebook/seamless-m4t-v2-large"
 stt_processor = AutoProcessor.from_pretrained(STT_MODEL)
 stt_model = SeamlessM4Tv2ForSpeechToText.from_pretrained(STT_MODEL)
-stt_model = stt_model.to(device)
-stt_model.eval()
+stt_model = stt_model.to(device).eval()
 print("✅ SeamlessM4T v2 Large loaded!")
 
-# NLLB-200 for Translation
+# 2. NLLB-200 for Translation
 print("📥 Loading NLLB-200...")
 NLLB_MODEL = "facebook/nllb-200-distilled-600M"
 nllb_tokenizer = AutoTokenizer.from_pretrained(NLLB_MODEL)
 nllb_model = AutoModelForSeq2SeqLM.from_pretrained(NLLB_MODEL)
-nllb_model = nllb_model.to(device)
-nllb_model.eval()
+nllb_model = nllb_model.to(device).eval()
 print("✅ NLLB-200 loaded!")
 
+# 3. SeamlessExpressive for Expressive Speech Translation
+print("📥 Loading SeamlessExpressive...")
+EXPRESSIVE_MODEL = "facebook/seamless-expressive"
+try:
+    exp_processor = AutoProcessor.from_pretrained(EXPRESSIVE_MODEL)
+    exp_model = SeamlessM4Tv2Model.from_pretrained(EXPRESSIVE_MODEL)
+    exp_model = exp_model.to(device).eval()
+    EXPRESSIVE_AVAILABLE = True
+    print("✅ SeamlessExpressive loaded!")
+except Exception as e:
+    EXPRESSIVE_AVAILABLE = False
+    print(f"⚠️ SeamlessExpressive not available: {e}")
+
 print("🎉 All models ready!")
 
 # ============================================================
@@ -55,11 +68,12 @@ STT_LANGS = {
     "English": "eng", "French": "fra", "Arabic": "arb", "Spanish": "spa",
     "German": "deu", "Italian": "ita", "Portuguese": "por", "Chinese": "cmn",
     "Japanese": "jpn", "Korean": "kor", "Russian": "rus", "Turkish": "tur",
-    "Dutch": "nld", "Hindi": "hin",
 }
 
+EXPRESSIVE_LANGS = ["English", "French", "German", "Spanish", "Italian", "Chinese"]
+
 # ============================================================
-# STT Function
+# STT Function (SeamlessM4T v2 Large)
 # ============================================================
 
 def stt(audio, src_lang):
@@ -97,7 +111,7 @@ def stt(audio, src_lang):
         return f"Error: {str(e)}"
 
 # ============================================================
-# Translation Function
+# Translation Function (NLLB-200)
 # ============================================================
 
 def translate(text, src_lang, tgt_lang):
@@ -126,13 +140,61 @@ def translate(text, src_lang, tgt_lang):
     except Exception as e:
         return f"Error: {str(e)}"
 
+# ============================================================
+# Expressive Speech Translation (SeamlessExpressive)
+# ============================================================
+
+def expressive_translate(audio, src_lang, tgt_lang):
+    """Expressive Speech-to-Speech Translation"""
+    if not EXPRESSIVE_AVAILABLE:
+        return None, "SeamlessExpressive not available"
+    
+    if audio is None:
+        return None, "No audio provided"
+    
+    try:
+        if isinstance(audio, tuple):
+            sample_rate, audio_data = audio
+            audio_data = audio_data.astype(np.float32)
+            if np.abs(audio_data).max() > 1.0:
+                audio_data = audio_data / 32768.0
+        else:
+            return None, "Invalid audio format"
+        
+        src_code = STT_LANGS.get(src_lang, "eng")
+        tgt_code = STT_LANGS.get(tgt_lang, "fra")
+        
+        inputs = exp_processor(
+            audios=audio_data,
+            sampling_rate=sample_rate,
+            return_tensors="pt"
+        ).to(device)
+        
+        with torch.no_grad():
+            output = exp_model.generate(
+                **inputs,
+                tgt_lang=tgt_code,
+                return_intermediate_token_ids=True
+            )
+        
+        # Get audio output
+        audio_output = output.audio_sequences[0].cpu().numpy()
+        
+        # Get text
+        text = exp_processor.decode(output.sequences[0].tolist(), skip_special_tokens=True)
+        
+        return (16000, audio_output), text
+        
+    except Exception as e:
+        return None, f"Error: {str(e)}"
+
 # ============================================================
 # Gradio Interface
 # ============================================================
 
 with gr.Blocks(title="STTR API", theme=gr.themes.Soft()) as demo:
     gr.Markdown("# 🌍 STTR - Speech & Translation API")
-    gr.Markdown("**SeamlessM4T v2 Large** + **NLLB-200** (200 languages + Darija!)")
+    gr.Markdown("**Meta AI Models:** SeamlessM4T v2 Large + NLLB-200 + SeamlessExpressive")
     
     with gr.Tab("🎤 Speech-to-Text"):
         stt_audio = gr.Audio(label="Audio", type="numpy")
@@ -149,5 +211,16 @@ with gr.Blocks(title="STTR API", theme=gr.themes.Soft()) as demo:
         trans_output = gr.Textbox(label="Translation", lines=3)
         trans_btn = gr.Button("🌍 Translate", variant="primary")
         trans_btn.click(translate, [trans_text, trans_src, trans_tgt], trans_output, api_name="translate")
+    
+    with gr.Tab("🎭 Expressive (S2S)"):
+        gr.Markdown("**SeamlessExpressive** - Preserves tone, emotion & style!")
+        exp_audio = gr.Audio(label="Input Audio", type="numpy")
+        with gr.Row():
+            exp_src = gr.Dropdown(EXPRESSIVE_LANGS, label="From", value="English")
+            exp_tgt = gr.Dropdown(EXPRESSIVE_LANGS, label="To", value="French")
+        exp_output_audio = gr.Audio(label="Translated Audio")
+        exp_output_text = gr.Textbox(label="Translated Text")
+        exp_btn = gr.Button("🎭 Translate with Expression", variant="primary")
+        exp_btn.click(expressive_translate, [exp_audio, exp_src, exp_tgt], [exp_output_audio, exp_output_text], api_name="expressive")
 
 demo.launch()