import whisper
import gradio as gr
import os
from pydub import AudioSegment
from transformers import BlenderbotTokenizer, BlenderbotForConditionalGeneration
from TTS.api import TTS
import torch

mname = "facebook/blenderbot-400M-distill"
chat_model = BlenderbotForConditionalGeneration.from_pretrained(mname)
tokenizer = BlenderbotTokenizer.from_pretrained(mname)

device = 'cuda' if torch.cuda.is_available() else 'cpu'

chat_model = chat_model.to(device)
SR_model = whisper.load_model("base")
SR_model = SR_model.to(device)

tts_model_name = TTS.list_models()[11]
tts = TTS(tts_model_name)

def querybot(audio_file):
    AudioSegment.from_wav(audio_file).export("file.mp3", format="mp3")
    UTTERANCE = SR_model.transcribe("file.mp3")["text"]

    inputs = tokenizer([UTTERANCE], return_tensors="pt").to(device)
    reply_ids = chat_model.generate(**inputs)
    response = tokenizer.batch_decode(reply_ids, skip_special_tokens=True)[0]
    tts.tts_to_file(text=response, file_path='tts.mp3')

    output = response

    return output, 'tts.mp3'


gr.Interface(
    fn=querybot, 
    inputs=gr.Audio(source="microphone", type="filepath"), 
    outputs=['text',gr.Audio(type='filepath')]).launch()