import gradio as gr
import json
import torch
import wavio
from tqdm import tqdm
from huggingface_hub import snapshot_download
from models import AudioDiffusion, DDPMScheduler
from audioldm.audio.stft import TacotronSTFT
from audioldm.variational_autoencoder import AutoencoderKL
from gradio import Markdown

class Tango:
    def __init__(self, name="declare-lab/tango", device="cuda:0"):
        
        path = snapshot_download(repo_id=name)
        
        vae_config = json.load(open("{}/vae_config.json".format(path)))
        stft_config = json.load(open("{}/stft_config.json".format(path)))
        main_config = json.load(open("{}/main_config.json".format(path)))
        
        self.vae = AutoencoderKL(**vae_config).to(device)
        self.stft = TacotronSTFT(**stft_config).to(device)
        self.model = AudioDiffusion(**main_config).to(device)
        
        vae_weights = torch.load("{}/pytorch_model_vae.bin".format(path), map_location=device)
        stft_weights = torch.load("{}/pytorch_model_stft.bin".format(path), map_location=device)
        main_weights = torch.load("{}/pytorch_model_main.bin".format(path), map_location=device)
        
        self.vae.load_state_dict(vae_weights)
        self.stft.load_state_dict(stft_weights)
        self.model.load_state_dict(main_weights)

        print ("Successfully loaded checkpoint from:", name)
        
        self.vae.eval()
        self.stft.eval()
        self.model.eval()
        
        self.scheduler = DDPMScheduler.from_pretrained(main_config["scheduler_name"], subfolder="scheduler")
        
    def chunks(self, lst, n):
        """ Yield successive n-sized chunks from a list. """
        for i in range(0, len(lst), n):
            yield lst[i:i + n]
        
    def generate(self, prompt, steps=100, guidance=3, samples=1, disable_progress=True):
        """ Genrate audio for a single prompt string. """
        with torch.no_grad():
            latents = self.model.inference([prompt], self.scheduler, steps, guidance, samples, disable_progress=disable_progress)
            mel = self.vae.decode_first_stage(latents)
            wave = self.vae.decode_to_waveform(mel)
        return wave[0]
    
    def generate_for_batch(self, prompts, steps=200, guidance=3, samples=1, batch_size=8, disable_progress=True):
        """ Genrate audio for a list of prompt strings. """
        outputs = []
        for k in tqdm(range(0, len(prompts), batch_size)):
            batch = prompts[k: k+batch_size]
            with torch.no_grad():
                latents = self.model.inference(batch, self.scheduler, steps, guidance, samples, disable_progress=disable_progress)
                mel = self.vae.decode_first_stage(latents)
                wave = self.vae.decode_to_waveform(mel)
                outputs += [item for item in wave]
        if samples == 1:
            return outputs
        else:
            return list(self.chunks(outputs, samples))

# Initialize TANGO
if torch.cuda.is_available():
    tango = Tango()
else:
    tango = Tango(device="cpu")

def gradio_generate(prompt, steps, guidance):
    output_wave = tango.generate(prompt, steps, guidance)
    output_filename = "temp_output.wav"
    wavio.write(output_filename, output_wave, rate=16000, sampwidth=2)
    
    return output_filename

description_text = "Generate audio using TANGO by providing a text prompt. \
\n\nLimitations: TANGO is trained on the small AudioCaps dataset so it may not generate good audio \
samples related to concepts that it has not seen in training (e.g. singing). For the same reason, TANGO \
is not always able to finely control its generations over textual control prompts. For example, \
the generations from TANGO for prompts Chopping tomatoes on a wooden table and Chopping potatoes \
on a metal table are very similar. \
\n\nWe are currently training another version of TANGO on larger datasets to enhance its generalization, \
compositional and controllable generation ability."

# Gradio input and output components
input_text = gr.inputs.Textbox(lines=2, label="Prompt")
output_audio = gr.outputs.Audio(label="Generated Audio", type="filepath")
denoising_steps = gr.Slider(minimum=100, maximum=200, value=100, step=1, label="Steps", interactive=True)
guidance_scale = gr.Slider(minimum=1, maximum=10, value=3, step=0.1, label="Guidance Scale", interactive=True)

# Gradio interface
gr_interface = gr.Interface(
    fn=gradio_generate,
    inputs=[input_text, denoising_steps, guidance_scale],
    outputs=[output_audio],
    title="TANGO: Text to Audio using Instruction-Guided Diffusion",
    description=description_text,
    allow_flagging=False,
    examples=[
        ["An audience cheering and clapping"],
        ["Rolling thunder with lightning strikes"],
        ["Gentle water stream, birds chirping and sudden gun shot"],
        ["A car engine revving"],
        ["A dog barking"],
        ["A cat meowing"],
        ["Wooden table tapping sound while water pouring"],
        ["Emergency sirens wailing"],
        ["two gunshots followed by birds flying away while chirping"],
        ["Whistling with birds chirping"],
        ["A person snoring"],
        ["Motor vehicles are driving with loud engines and a person whistles"],
        ["People cheering in a stadium while thunder and lightning strikes"],
        ["A helicopter is in flight"],
        ["A dog barking and a man talking and a racing car passes by"],
    ],
    cache_examples=False,
)

# Launch Gradio app
gr_interface.launch()