import gradio as gr
from diffusers import AudioLDMControlNetPipeline, ControlNetModel
import os
from pretty_midi import PrettyMIDI
from tempfile import _TemporaryFileWrapper
import torch
import torchaudio

if torch.cuda.is_available():
    device = "cuda"
    torch_dtype = torch.float16
else:
    device = "cpu"
    torch_dtype = torch.float32

controlnet = ControlNetModel.from_pretrained(
    "lauraibnz/midi-audioldm", torch_dtype=torch_dtype)
pipe = AudioLDMControlNetPipeline.from_pretrained(
    "cvssp/audioldm-m-full", controlnet=controlnet, torch_dtype=torch_dtype)
pipe = pipe.to(device)
generator = torch.Generator(device)


def predict(midi_file=None, prompt="", negative_prompt="", audio_length_in_s=5, random_seed=0, controlnet_conditioning_scale=1, num_inference_steps=20, guess_mode=False):
    if isinstance(midi_file, _TemporaryFileWrapper):
        midi_file = midi_file.name
    midi = PrettyMIDI(midi_file)
    audio = pipe(
        prompt,
        negative_prompt=negative_prompt,
        midi=midi,
        audio_length_in_s=audio_length_in_s,
        num_inference_steps=num_inference_steps,
        controlnet_conditioning_scale=float(controlnet_conditioning_scale),
        guess_mode=guess_mode,
        generator=generator.manual_seed(int(random_seed)),
    )
    return (16000, audio.audios.T)

with gr.Blocks(title="🎹 MIDI-AudioLDM", theme=gr.themes.Base(text_size=gr.themes.sizes.text_md, font=[gr.themes.GoogleFont("Nunito Sans")])) as demo:
    gr.HTML(
        """
        <h1 align="center"; size="16">🎹 MIDI-AudioLDM </h1>
        """)
    gr.Markdown(
            """
            MIDI-AudioLDM is a MIDI-conditioned text-to-audio model based on the project [AudioLDM](https://huggingface.co/spaces/haoheliu/audioldm-text-to-audio-generation). The model has been conditioned using the ControlNet architecture and has been developed within Hugging Face’s [🧨 Diffusers](https://huggingface.co/docs/diffusers/) framework. Once trained, MIDI-AudioLDM accepts a MIDI file and a text prompt as input and returns an audio file, which is an interpretation of the MIDI based on the given text description. This enables detailed control over different musical aspects such as notes, mood and timbre.
            """)
    with gr.Row():
        with gr.Column(variant='panel'):
            midi = gr.File(label="midi file", file_types=[".mid"], info="Load the MIDI file that you want to use as conditioning.")
            prompt = gr.Textbox(label="prompt", info="Enter a descriptive text prompt.")
        with gr.Column(variant='panel'):
            audio = gr.Audio(label="audio")
    with gr.Accordion("Advanced Settings", open=False):
        neg_prompt = gr.Textbox(label="negative prompt", info="Enter a negative prompt not to guide the audio generation.")
        duration = gr.Slider(0, 30, value=5, step=5, label="duration (seconds)", info="Modify the duration of the output audio file.")
        seed = gr.Number(value=42, label="seed", info="Change the random seed for a different generation result.")
        cond = gr.Slider(0.0, 1.0, value=1.0, step=0.1, label="conditioning scale", info="Enter a value between 0 and 1. The larger the more it will take the conditioning into account.")
        inf = gr.Slider(0, 50, value=20, step=0.1, label="inference steps", info="Edit the number of denoising steps. More inference steps usually leads to better but slower results.")
        guess = gr.Checkbox(label="guess mode", info="If true, the model will try to recognize the content of the conditioning without the need of a text prompt.")
    btn = gr.Button("Generate")
    btn.click(predict, inputs=[midi, prompt, neg_prompt, duration, seed, cond, inf, guess], outputs=[audio])
    gr.Examples(examples=[["S00.mid", "piano", "", 10, 25, 1.0, 20, False], ["S00.mid", "violin", "", 10, 25, 1.0, 20, False], ["S00.mid", "woman singing", "", 10, 25, 0.8, 20, False]], inputs=[midi, prompt, neg_prompt, duration, seed, cond, inf, guess], fn=predict, outputs=audio, cache_examples=True)
    
demo.launch()