import gradio as gr
import os
import torch
import soundfile as sf
from transformers import Wav2Vec2Processor, Wav2Vec2ForCTC
from TTS.tts.configs.xtts_config import XttsConfig
from TTS.tts.models.xtts import Xtts
import librosa
import string

os.makedirs("output_audio", exist_ok=True)

# Initialize ASR model
asr_processor = Wav2Vec2Processor.from_pretrained("facebook/wav2vec2-large-960h")
asr_model = Wav2Vec2ForCTC.from_pretrained("facebook/wav2vec2-large-960h")
asr_model.eval()

# Initialize TTS model
tts_config_path = "/path/to/xtts/config.json"
tts_checkpoint_dir = "/path/to/xtts/"
speaker_wav_path = "/path/to/target/speaker.wav"  # Update with actual speaker wav path for cloning voice

tts_config = XttsConfig()
tts_config.load_json(tts_config_path)
tts_model = Xtts.init_from_config(tts_config)
tts_model.load_checkpoint(tts_config, checkpoint_dir=tts_checkpoint_dir, eval=True)
tts_model.cuda()

def resample_audio(input_audio_path, target_sr):
    waveform, sr = sf.read(input_audio_path)
    if sr != target_sr:
        waveform = librosa.resample(waveform, orig_sr=sr, target_sr=target_sr)
    return waveform

def speech_to_text(input_audio_path):
    waveform = resample_audio(input_audio_path, 16000)
    input_values = asr_processor(waveform, sampling_rate=16000, return_tensors="pt").input_values
    with torch.no_grad():
        logits = asr_model(input_values).logits
    predicted_ids = torch.argmax(logits, dim=-1)
    transcription = asr_processor.batch_decode(predicted_ids)[0]
    return transcription.strip()

def text_to_speech(text, output_path="output_audio/output.wav"):
    if not text.strip():
        return "The text input is empty, please provide a valid string."

    outputs = tts_model.synthesize(
        text,
        tts_config,
        speaker_wav=speaker_wav_path,
        gpt_cond_len=3,
        language="en"
    )
    waveform = outputs['waveform'].squeeze().cpu().numpy()
    sf.write(output_path, waveform, 22050)
    return output_path

def speech_to_speech(input_audio, text_input=None):
    if text_input is None:
        text_input = speech_to_text(input_audio)
    synthesized_speech_path = text_to_speech(text_input)
    return synthesized_speech_path

iface = gr.Interface(
    fn=speech_to_speech,
    inputs=[gr.Audio(type="filepath", label="Input Audio"), gr.Textbox(label="Text Input", optional=True)],
    outputs=gr.Audio(label="Synthesized Speech"),
    title="Speech-to-Speech Application",
    description="Converts speech to text and then back to speech, ensuring the output audio is of high quality."
)

iface.launch(share=True)