app.py

import torch
import gradio as gr
from faster_whisper import WhisperModel
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
from pydub import AudioSegment
import yt_dlp as youtube_dl
import tempfile
from transformers.pipelines.audio_utils import ffmpeg_read
from gradio.components import Audio, Dropdown, Radio, Textbox
import os
import numpy as np
import soundfile as sf
os.environ["TOKENIZERS_PARALLELISM"] = "false"


# Paramètres
FILE_LIMIT_MB = 1000
YT_LENGTH_LIMIT_S = 3600  # Limite de 1 heure pour les vidéos YouTube

# Charger les codes de langue
from flores200_codes import flores_codes

# Fonction pour déterminer le device
def set_device():
    return torch.device("cuda" if torch.cuda.is_available() else "cpu")

device = set_device()


# Charger les modèles une seule fois
model_dict = {}
def load_models():
    global model_dict
    if not model_dict:
        model_name_dict = {
            #'nllb-distilled-1.3B': 'facebook/nllb-200-distilled-1.3B',
            'nllb-distilled-600M': 'facebook/nllb-200-distilled-600M',
            #'nllb-1.3B': 'facebook/nllb-200-1.3B',
            #'nllb-distilled-1.3B': 'facebook/nllb-200-distilled-1.3B',
            #'nllb-3.3B': 'facebook/nllb-200-3.3B',
            # 'nllb-distilled-600M': 'facebook/nllb-200-distilled-600M',
        }
        for call_name, real_name in model_name_dict.items():
            model = AutoModelForSeq2SeqLM.from_pretrained(real_name)
            tokenizer = AutoTokenizer.from_pretrained(real_name)
            model_dict[call_name+'_model'] = model
            model_dict[call_name+'_tokenizer'] = tokenizer

load_models()

model_size = "large-v2"
model = WhisperModel(model_size)


# Fonction pour la transcription
def transcribe_audio(audio_file):
    # model_size = "large-v2"
    # model = WhisperModel(model_size)
    # model = WhisperModel(model_size, device=device, compute_type="int8")
    global model
    segments, _ = model.transcribe(audio_file, beam_size=1)
    transcriptions = [("[%.2fs -> %.2fs]" % (seg.start, seg.end), seg.text) for seg in segments]
    return transcriptions


# Fonction pour la traduction
def traduction(text, source_lang, target_lang):
    # Vérifier si les codes de langue sont dans flores_codes
    if source_lang not in flores_codes or target_lang not in flores_codes:
        print(f"Code de langue non trouvé : {source_lang} ou {target_lang}")
        return ""

    src_code = flores_codes[source_lang]
    tgt_code = flores_codes[target_lang]

    model_name = "nllb-distilled-600M"
    model = model_dict[model_name + "_model"]
    tokenizer = model_dict[model_name + "_tokenizer"]
    translator = pipeline("translation", model=model, tokenizer=tokenizer)

    return translator(text, src_lang=src_code, tgt_lang=tgt_code)[0]["translation_text"]


# Fonction principale
def full_transcription_and_translation(audio_input, source_lang, target_lang):
    # Si audio_input est une URL
    if isinstance(audio_input, str) and audio_input.startswith("http"):
        audio_file = download_yt_audio(audio_input)
    # Si audio_input est un dictionnaire contenant des données audio
    elif isinstance(audio_input, dict) and "array" in audio_input and "sampling_rate" in audio_input:
        audio_array = audio_input["array"]
        sampling_rate = audio_input["sampling_rate"]
        # Écrire le tableau NumPy dans un fichier temporaire WAV
        with tempfile.NamedTemporaryFile(delete=False, suffix='.wav') as f:
            sf.write(f, audio_array, sampling_rate)
            audio_file = f.name
    else:
        # Supposons que c'est un chemin de fichier
        audio_file = audio_input

    transcriptions = transcribe_audio(audio_file)
    translations = [(timestamp, traduction(text, source_lang, target_lang)) for timestamp, text in transcriptions]

    # Supprimez le fichier temporaire s'il a été créé
    if isinstance(audio_input, dict):
        os.remove(audio_file)

    return transcriptions, translations

# Téléchargement audio YouTube
"""def download_yt_audio(yt_url):
    with tempfile.NamedTemporaryFile(suffix='.mp3') as f:
        ydl_opts = {
            'format': 'bestaudio/best',
            'outtmpl': f.name,
            'postprocessors': [{
                'key': 'FFmpegExtractAudio',
                'preferredcodec': 'mp3',
                'preferredquality': '192',
            }],
        }
        with youtube_dl.YoutubeDL(ydl_opts) as ydl:
            ydl.download([yt_url])
        return f.name"""

lang_codes = list(flores_codes.keys())

# Interface Gradio
def gradio_interface(audio_file, source_lang, target_lang):
    if audio_file.startswith("http"):
        audio_file = download_yt_audio(audio_file)
    transcriptions, translations = full_transcription_and_translation(audio_file, source_lang, target_lang)
    transcribed_text = '\n'.join([f"{timestamp}: {text}" for timestamp, text in transcriptions])
    translated_text = '\n'.join([f"{timestamp}: {text}" for timestamp, text in translations])
    return transcribed_text, translated_text


def _return_yt_html_embed(yt_url):
    video_id = yt_url.split("?v=")[-1]
    HTML_str = (
        f'<center> <iframe width="500" height="320" src="https://www.youtube.com/embed/{video_id}"> </iframe>'
        " </center>"
    )
    return HTML_str

def download_yt_audio(yt_url, filename):
    info_loader = youtube_dl.YoutubeDL()

    try:
        info = info_loader.extract_info(yt_url, download=False)
    except youtube_dl.utils.DownloadError as err:
        raise gr.Error(str(err))

    file_length = info["duration_string"]
    file_h_m_s = file_length.split(":")
    file_h_m_s = [int(sub_length) for sub_length in file_h_m_s]

    if len(file_h_m_s) == 1:
        file_h_m_s.insert(0, 0)
    if len(file_h_m_s) == 2:
        file_h_m_s.insert(0, 0)
    file_length_s = file_h_m_s[0] * 3600 + file_h_m_s[1] * 60 + file_h_m_s[2]

    if file_length_s > YT_LENGTH_LIMIT_S:
        yt_length_limit_hms = time.strftime("%HH:%MM:%SS", time.gmtime(YT_LENGTH_LIMIT_S))
        file_length_hms = time.strftime("%HH:%MM:%SS", time.gmtime(file_length_s))
        raise gr.Error(f"Maximum YouTube length is {yt_length_limit_hms}, got {file_length_hms} YouTube video.")

    ydl_opts = {"outtmpl": filename, "format": "worstvideo[ext=mp4]+bestaudio[ext=m4a]/best[ext=mp4]/best"}

    with youtube_dl.YoutubeDL(ydl_opts) as ydl:
        try:
            ydl.download([yt_url])
        except youtube_dl.utils.ExtractorError as err:
            raise gr.Error(str(err))


def yt_transcribe(yt_url, task, max_filesize=75.0):
    html_embed_str = _return_yt_html_embed(yt_url)
    global model  # S'assurer que le modèle est accessible

    with tempfile.TemporaryDirectory() as tmpdirname:
        filepath = os.path.join(tmpdirname, "video.mp4")
        download_yt_audio(yt_url, filepath)
        with open(filepath, "rb") as f:
            inputs = f.read()

    inputs = ffmpeg_read(inputs, model.feature_extractor.sampling_rate)
    inputs = {"array": inputs, "sampling_rate": model.feature_extractor.sampling_rate}

    transcriptions, translations = full_transcription_and_translation(inputs, source_lang, target_lang)
    transcribed_text = '\n'.join([f"{timestamp}: {text}" for timestamp, text in transcriptions])
    translated_text = '\n'.join([f"{timestamp}: {text}" for timestamp, text in translations])
    return html_embed_str, transcribed_text, translated_text


# Interfaces
demo = gr.Blocks()

with demo:
    with gr.Tab("Microphone"):
        gr.Interface(
            fn=gradio_interface,
            inputs=[
                gr.Audio(sources=["microphone"], type="filepath"), 
                gr.Dropdown(lang_codes, value='French', label='Source Language'), 
                gr.Dropdown(lang_codes, value='English', label='Target Language')],
            outputs=[gr.Textbox(label="Transcribed Text"), gr.Textbox(label="Translated Text")]
        )

    with gr.Tab("Audio file"):
        gr.Interface(
            fn=gradio_interface,
            inputs=[
                gr.Audio(type="filepath", label="Audio file"), 
                gr.Dropdown(lang_codes, value='French', label='Source Language'), 
                gr.Dropdown(lang_codes, value='English', label='Target Language')],
            outputs=[gr.Textbox(label="Transcribed Text"), gr.Textbox(label="Translated Text")]
        )

    with gr.Tab("YouTube"):
        gr.Interface(
            fn=yt_transcribe,
            inputs=[
                gr.Textbox(lines=1, placeholder="Paste the URL to a YouTube video here", label="YouTube URL"), 
                gr.Dropdown(lang_codes, value='French', label='Source Language'), 
                gr.Dropdown(lang_codes, value='English', label='Target Language')
            ],
            outputs=["html", gr.Textbox(label="Transcribed Text"), gr.Textbox(label="Translated Text")]
        )

#with demo:
    #gr.TabbedInterface([mf_transcribe, file_transcribe, yt_transcribe], ["Microphone", "Audio file", "YouTube"])

demo.launch()