add api key text input

app.py

@@ -4,6 +4,10 @@ import numpy as np
 import pretty_midi
 import re
 import numpy as np
+import os
+import gradio as gr
+
+openai.api_key = os.environ.get("OPENAI_API_KEY")
 
 # sample data
 markdown_table_sample = """4
@@ -45,8 +49,7 @@ MIDI_NOTENUM = {
 }
 SR = 44100
 
-count = 0
-MAX_QUERY = 100
+MAX_QUERY = 5
 
 def convert_table_to_audio(markdown_table, resolution=8, bpm = 120.0):
     # convert table to array
@@ -92,11 +95,10 @@ def get_answer(question):
     )
     return response["choices"][0]["message"]["content"]
 
-def generate_rhythm(query):
-    global count 
-    count += 1
-    if count > MAX_QUERY:
-        return [None, "Now you can try up to %d times" % MAX_QUERY] 
+def generate_rhythm(query, state):
+    if state["gen_count"] > MAX_QUERY:
+        return [None, "You need to set your ChatGPT API Key to try more than %d times" % MAX_QUERY] 
+    state["gen_count"] = state["gen_count"] + 1
 
     # get respance from ChatGPT
     text_output = get_answer(query)
@@ -119,17 +121,25 @@ def generate_rhythm(query):
     return [(SR, audio_data), text_output] 
 # %%
 
-import gradio as gr
+def on_token_change(user_token):
+    openai.api_key = user_token or os.environ.get("OPENAI_API_KEY")
+
 
 with gr.Blocks() as demo:
+    state = gr.State({"gen_count": 0})
+
     gr.Markdown("Ask ChatGPT to generate rhythm patterns")
     with gr.Row():
-        inp = gr.Textbox(placeholder="Give me a Hiphop rhythm pattern with some reggae twist!")
+        with gr.Column():
+            inp = gr.Textbox(placeholder="Give me a Hiphop rhythm pattern with some reggae twist!")
+            btn = gr.Button("Generate")
         with gr.Column():
             out_audio = gr.Audio()
             out_text = gr.Textbox(placeholder="ChatGPT output")
-    btn = gr.Button("Generate")
-    btn.click(fn=generate_rhythm, inputs=inp, outputs=[out_audio, out_text])
+    gr.Markdown("Enter your own OpenAI API Key to try out more than 5 times. You can get it [here](https://platform.openai.com/account/api-keys).", elem_id="label")
+    user_token = gr.Textbox(placeholder="OpenAI API Key", type="password", show_label=False)
+    btn.click(fn=generate_rhythm, inputs=[inp, state], outputs=[out_audio, out_text])
+    user_token.change(on_token_change, inputs=[user_token], outputs=[])
 demo.launch()
 # demo = gr.Interface(
 #     fn=generate_rhythm,