import gradio as gr
import os
import numpy as np

def no_cpu_warning(video_file, dataset, lora, size):
    output = "./cached_examples/"
    output += "peft" if lora else "audiocraft"
    output += "_" + dataset + "_" + size + "_" + video_file[-7:-4]
    chosen_video = np.random.choice(['v', 'v1', 'v2'])
    output += "_" + chosen_video + ".mp4"

    if not os.path.exists(output):
        print(output)
        raise gr.Error("This combination of video and model has not been cached. Try one of the listed examples instead.")

    gr.Warning("This interface is running on CPU. Only cached generated examples are displayed here, no new music is generated. Please be patient, as the videos meight take a while to download.")
    print("Displaying video: " + output)
    return output

with gr.Blocks() as demo:
    gr.Markdown('<span style="color:yellow"> **WARNING**: Multiple cached videos are examples are displayed here. Please be patient, as the videos meight take a while to download.</span>')
    gr.Interface(fn=no_cpu_warning,
                         inputs=[
                                 gr.Video(value="./videos/n_5.mp4",
                                          label="Original Video",
                                          min_length=5,
                                          max_length=20,
                                          sources=[],
                                          show_download_button=True,
                                          include_audio=True, 
                                          mirror_webcam=False
                                          ),
                                 gr.Radio(["nature", "symmv"],
                                          value="nature",
                                          label="Available Models",
                                          info="Choose one of the datasets on which the models has been trained on. Nature is a dataset of calm and relaxing sounds, symmv contains charts music. The model will generate audio that is similar to the training data."
                                          ),
                                 gr.Radio([False, True],
                                          label="Use the LoRA version of the MusicGen Audio Decoder",
                                          value=False,
                                          info="If set to 'True' the MusicGen Audio Decoder models trained with LoRA "
                                               "(Low Rank Adaptation) are used. If set to 'False', the original "
                                               "MusicGen models are used instead."
                                               ),
                                 gr.Radio(["small", "medium", "large"],
                                          label="Model Size",
                                          value="large",
                                          info="Choose one of the available model sizes. This effects the enocder and the deocder model. The larger models are more likely produce "
                                               "results of higher audio quality, but also take more time to generate it."
                                               ),
                                ],

                         outputs=[gr.Video(label="Generated Result")],
                         examples=[
                             [os.path.abspath("./videos/n_1.mp4"), "nature", False, "large"],
                             [os.path.abspath("./videos/n_2.mp4"), "nature", False, "large"],
                             [os.path.abspath("./videos/n_3.mp4"), "nature", False, "large"],
                             [os.path.abspath("./videos/n_4.mp4"), "nature", False, "large"],
                             [os.path.abspath("./videos/n_5.mp4"), "nature", True, "large"],
                             [os.path.abspath("./videos/n_6.mp4"), "nature", True, "large"],
                             [os.path.abspath("./videos/n_7.mp4"), "nature", True, "large"],
                             [os.path.abspath("./videos/n_8.mp4"), "nature", True, "large"],
                             [os.path.abspath("./videos/s_1.mp4"), "symmv", False, "large"],
                             [os.path.abspath("./videos/s_2.mp4"), "symmv", False, "large"],
                             [os.path.abspath("./videos/s_3.mp4"), "symmv", False, "large"],
                             [os.path.abspath("./videos/s_4.mp4"), "symmv", False, "large"],
                             [os.path.abspath("./videos/s_5.mp4"), "symmv", True, "large"],
                             [os.path.abspath("./videos/s_6.mp4"), "symmv", True, "large"],
                             [os.path.abspath("./videos/s_7.mp4"), "symmv", True, "large"],
                             [os.path.abspath("./videos/s_8.mp4"), "symmv", True, "large"],
                           ],
                         examples_per_page=20,
                         cache_examples=False,
                         )

if __name__ == "__main__":
    demo.launch(
        share=False
    )