import gradio as gr
import openai
from t2a import text_to_audio
import joblib
from sentence_transformers import SentenceTransformer
import numpy as np
import os

reg = joblib.load('text_reg.joblib')
model = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')
finetune = "davinci:ft-personal:autodrummer-v5-2022-11-04-22-34-07"

def get_note_text(prompt):
    prompt = prompt + " ->"
    # get completion from finetune
    response = openai.Completion.create(
        engine=finetune,
        prompt=prompt,
        temperature=0.5,
        max_tokens=200,
        top_p=1,
        frequency_penalty=0,
        presence_penalty=0,
        stop=["###"]
    )
    return response.choices[0].text.strip()

def get_drummer_output(prompt):
    openai.api_key = os.environ['key']
    note_text = get_note_text(prompt)
    # note_text = note_text + " " + note_text
    prompt_enc = model.encode([prompt])
    bpm = int(reg.predict(prompt_enc)[0]) + 20
    print(bpm, "bpm", "notes are", note_text)
    audio = text_to_audio(note_text, bpm)
    audio = np.array(audio.get_array_of_samples(), dtype=np.float32)
    return (96000, audio)

iface = gr.Interface(
    fn=get_drummer_output,
    inputs="text",
    examples=[
        "hiphop groove 808",
        "rock metal",
        "disco funk",
    ],
    outputs="audio",
    title='Autodrummer',
    description="Stable Diffusion for drum beats. Type in a genre and some descriptors (e.g., 'hiphop groove 808') to the prompt box and get a drum beat in that genre"
)
iface.launch()