import gradio as gr
import os
import numpy as np

def no_cpu_warning(video_file, dataset, lora, size):
    gr.Warning("This interface is running on CPU. Only cached examples are displayed here, no new music can be generated.")

    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. Please use the other model or try one of the listed examples instead.")

    print("Displaying video: " + output)
    return output

interface = gr.Interface(fn=no_cpu_warning,
                         inputs=[
                                 gr.Video(value="videos/n_5.mp4",
                                          label="Video Input",
                                          min_length=5,
                                          max_length=20,
                                          sources=['upload'],
                                          show_download_button=True,
                                          include_audio=True),
                                 gr.Radio(["nature", "symmv"],
                                          value="nature",
                                          label="Available Models",
                                          info="Choose one of the available Datasets on which the models has been trained on."),
                                 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. The larger models are more likely produce "
                                               "results of higher audio quality, but also take more time to generate it."),
                                ],

                         outputs=[gr.Video(label="video output")],
                         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"],
                           ],
                         cache_examples=False,
                         )

if __name__ == "__main__":
    interface.launch(
        share=True
    )