import streamlit as st
import torch
import tempfile
import os
import librosa
import numpy as np
from transformers import WhisperProcessor, WhisperForConditionalGeneration
from peft import PeftModel
import psutil

# Configuration de l'interface Streamlit
st.title("🔊 Transcription Audio avec Whisper Fine-tuné (LoRA)")
st.write("Upload un fichier audio et laisse ton modèle fine-tuné faire le travail !")

# 🔹 Charger le modèle Whisper Large et appliquer l’adaptateur LoRA
@st.cache_resource  # Permet de ne charger qu'une seule fois le modèle
def load_model():
    base_model_name = "openai/whisper-large"  # Modèle de base
    adapter_model_name = "SimpleFrog/whisper_finetuned"  # Adaptateur LoRA

    # Récupérer le token Hugging Face depuis une variable d'environnement
    hf_token = os.environ.get("HF_TOKEN")

    if not hf_token:
        raise ValueError("Aucun token Hugging Face trouvé dans l'environnement. Ajoute 'HF_TOKEN'.")


    # Charger le modèle de base
    model = WhisperForConditionalGeneration.from_pretrained(base_model_name)

    # Charger l'adaptateur LoRA (avec le token car modèle privé)
    model = PeftModel.from_pretrained(model, adapter_model_name, token=hf_token)

    # Charger le processeur audio
    processor = WhisperProcessor.from_pretrained(base_model_name)

    model.eval()  # Mode évaluation
    return processor, model

processor, model = load_model()

memory_used = psutil.Process(os.getpid()).memory_info().rss / (1024 * 1024)  # Convertir en MB
st.write(f"🖥️ Mémoire utilisée par le modèle : {memory_used:.2f} MB")

# Vérifier que les poids LoRA sont bien appliqués
if hasattr(model, "peft_config"):
    st.write("✅ Adaptateur LoRA chargé avec succès !")
    st.write("📂 Couches LoRA appliquées :", model.peft_config)
else:
    st.write("❌ Aucun adaptateur LoRA détecté, le modèle utilisé est Whisper Large standard.")

# 🔹 Upload d'un fichier audio
uploaded_file = st.file_uploader("Upload un fichier audio", type=["mp3", "wav", "m4a"])

if uploaded_file is not None:
    # Sauvegarder temporairement l'audio
    with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as temp_audio:
        temp_audio.write(uploaded_file.read())
        temp_audio_path = temp_audio.name

    # Charger et traiter l'audio
    st.write("📄 **Transcription en cours...**")

    # 🔹 Charger l'audio et convertir en waveform
    audio, sr = librosa.load(temp_audio_path, sr=16000)  # Whisper attend du 16kHz
    audio = np.expand_dims(audio, axis=0)  # Ajouter une dimension batch

    # 🔹 Préparer les entrées pour Whisper
    inputs = processor(audio, sampling_rate=16000, return_tensors="pt", language="fr", task="transcribe")

    # 🔹 Générer la transcription
    with torch.no_grad():
        predicted_ids = model.generate(input_features=inputs.input_features)

    # 🔹 Décoder la sortie
    transcription = processor.batch_decode(predicted_ids, skip_special_tokens=True)[0]

    # 🔹 Afficher la transcription
    st.subheader("📝 Transcription :")
    st.text_area("", transcription, height=200)

    # Supprimer le fichier temporaire après l'affichage
    os.remove(temp_audio_path)

st.write("🔹 Modèle utilisé :", "Whisper Large + Adaptateur LoRA (SimpleFrog/whisper_finetuned)")