Create asr-tts_service.py

asr-tts_service.py

@@ -0,0 +1,126 @@
+import base64
+from flask import Flask, request, jsonify
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
+import torch
+import soundfile as sf
+import numpy as np
+import io
+import re
+
+# Parler-TTS imports
+from parler_tts import ParlerTTSForConditionalGeneration
+
+# Flask App
+app = Flask(__name__)
+
+# ASR Wolof
+asr = pipeline("automatic-speech-recognition", model="bilalfaye/wav2vec2-large-mms-1b-wolof")
+
+# Translation Wolof <-> French
+model_name = "bilalfaye/nllb-200-distilled-600M-wo-fr-en"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+translation_model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
+
+fr_trans_model_name = "bilalfaye/nllb-200-distilled-600M-wo-fr-en"
+tokenizer_fr_trans = AutoTokenizer.from_pretrained(fr_trans_model_name, use_fast=False)
+fr_trans_model = AutoModelForSeq2SeqLM.from_pretrained(fr_trans_model_name)
+
+# Parler-TTS Model (Wolof)
+device = "cuda" if torch.cuda.is_available() else "cpu"
+tts_model = ParlerTTSForConditionalGeneration.from_pretrained("CONCREE/Adia_TTS").to(device)
+tts_tokenizer = AutoTokenizer.from_pretrained("CONCREE/Adia_TTS")
+tts_description = "A professional, clear and composed voice, perfect for formal presentations"
+
+# Helpers
+def wolofToFrench(wolof_text):
+    tokenizer.src_lang = "wol_Latn"
+    inputs = tokenizer(wolof_text, return_tensors="pt", padding=True)
+    forced_bos = tokenizer.convert_tokens_to_ids("fra_Latn")
+    translated_tokens = translation_model.generate(**inputs, forced_bos_token_id=forced_bos, max_new_tokens=200)
+    return tokenizer.batch_decode(translated_tokens, skip_special_tokens=True)[0]
+
+def frenchToWolof(fr_text):
+    tokenizer_fr_trans.src_lang = "fra_Latn"
+    inputs = tokenizer_fr_trans(fr_text, return_tensors="pt", padding=True)
+    forced_bos = tokenizer_fr_trans.convert_tokens_to_ids("wol_Latn")
+    translated_tokens = fr_trans_model.generate(**inputs, forced_bos_token_id=forced_bos, max_new_tokens=200)
+    return tokenizer_fr_trans.batch_decode(translated_tokens, skip_special_tokens=True)[0]
+
+def convert_digits_in_text(text):
+    # Exemple simple : remplacer 0 par "zéro", 1 par "un", etc.
+    digits_map = {"0":"zéro","1":"un","2":"deux","3":"trois","4":"quatre","5":"cinq","6":"six","7":"sept","8":"huit","9":"neuf"}
+    for k,v in digits_map.items():
+        text = text.replace(k, v)
+    return text
+
+def split_text(text, max_chars=170):
+    sentences = re.split(r'(?<=[.!?]) +', text)
+    chunks, current = [], ""
+    for s in sentences:
+        if len(current) + len(s) < max_chars:
+            current += " " + s if current else s
+        else:
+            chunks.append(current.strip())
+            current = s
+    if current:
+        chunks.append(current.strip())
+    return chunks
+
+
+# Routes
+@app.route("/", methods=["GET"])
+def racine():
+    return "Flask Asr-Tts Service is running!"
+
+@app.route("/accueil", methods=["GET"])
+def accueil():
+    return "Flask Asr-Tts accueil endpoint is working!"
+
+@app.route("/transcribe", methods=["POST"])
+def transcribe():
+    if "file" not in request.files:
+        return jsonify({"error": "Aucun fichier audio trouvé"}), 400
+    audio_file = request.files["file"]
+    data, samplerate = sf.read(audio_file)
+    text = asr(np.array(data))["text"]
+    translated = wolofToFrench(text)
+    return translated or "Bonjour Adama"
+
+# TTS Route 
+@app.route("/tts", methods=["POST"])
+def tts_route():
+    payload = request.get_json()
+    if not payload or "text" not in payload:
+        return jsonify({"error": "Champ 'text' manquant"}), 400
+
+    text_fr = payload["text"]
+    text_wolof = frenchToWolof(text_fr)
+    text_wolof = convert_digits_in_text(text_wolof)
+    chunks = split_text(text_wolof)
+
+    print("TTS chunks:", chunks)
+
+    audio_segments = []
+
+    tts_input_ids = tts_tokenizer(tts_description, return_tensors="pt").input_ids.to(device)
+
+    with torch.no_grad():
+        for chunk in chunks:
+            prompt_ids = tts_tokenizer(chunk, return_tensors="pt").input_ids.to(device)
+            audio_tensor = tts_model.generate(input_ids=tts_input_ids, prompt_input_ids=prompt_ids)
+            audio_segments.append(audio_tensor.cpu().numpy().squeeze())
+
+    final_audio = np.concatenate(audio_segments)
+    
+    buffer = io.BytesIO()
+    sf.write(buffer, final_audio, tts_model.config.sampling_rate, format="WAV")
+    buffer.seek(0)
+
+    audio_b64 = f"data:audio/wav;base64,{base64.b64encode(buffer.read()).decode('utf-8')}"
+
+    return jsonify({"audio": audio_b64})
+
+# Run Flask
+
+if __name__ == "__main__":
+    app.run(debug=False, host='0.0.0.0', port=7860, use_reloader=False)