Update app.py

app.py

@@ -1,6 +1,8 @@
 import random
 import os
 import time
+import requests
+
 from queue import Queue
 from threading import Thread
 
@@ -91,6 +93,8 @@ def postprocess(txt, path):
 
 
 
+
+
 with gr.Blocks() as demo:
     chatbot_box = gr.Chatbot(type="messages", render_markdown=False, sanitize_html=False)
     prefix_box = gr.TextArea(value="Twinkle Twinkle Little Star", label="Score title / text prefix")
@@ -104,9 +108,21 @@ with gr.Blocks() as demo:
     audio_box = gr.Audio()
     midi_box = gr.File()
     piano_roll_box = gr.Image()
+    server_box = gr.Dropdown(
+        choices=["Huggingface ZeroGPU", "CPU"], 
+        label="GPU Server",
+    )
+    gr.Markdown('''
+ZeroGPU comes with a time limit currently:
+- 3 minutes (not logged in)
+- 5 minutes (logged in)
+- 25 minutes (Pro user)
+
+CPUs will be slower but there is no time limit.
+'''.strip())
     example_box = gr.Examples(
         [
-            [example_prefix],
+            # [example_prefix],
             ["Twinkle Twinkle Little Star"], ["Twinkle Twinkle Little Star (Minor Key Version)"],
             ["The Entertainer - Scott Joplin (Piano Solo)"], ["Clair de Lune – Debussy"], ["Nocturne | Frederic Chopin"],
             ["Fugue I in C major, BWV 846"], ["Beethoven Symphony No. 7 (2nd movement) Piano solo"], 
@@ -120,27 +136,47 @@ with gr.Blocks() as demo:
     def user_fn(user_message, history: list):
         return "", history + [{"role": "user", "content": user_message}]
     
-    def generate_fn(history: list):
+    def get_last(history: list):
+        if len(history) == 0:
+            raise gr.Error('''No messages to read yet. Try the "Generate" button first!''')
+        return history[-1]["content"]
+    
+    def generate_fn(history, server):
          # continue from user input
-        prefix = history[-1]["content"]  
+        prefix = get_last(history)
         # prevent the model from continuing user's score title
         if prefix != '' and '\n' not in prefix:
-            # prefix is a single line --> prefix is the score title
-            prefix += '\n'
+            # prefix is a single line => prefix is the score title
+            # add '\n' to prevent model from continuing the title
+            prefix += '\n'  
             
         history.append({"role": "assistant", "content": ""})
         # history[-1]["content"] += "Generating with the given prefix...\n"
-        for history in model_fn(prefix, history):
+        for history in model_fn(prefix, history, server):
             yield history
 
-    def continue_fn(history: list):
+    def continue_fn(history, server):
          # continue from the last model output
         prefix = history[-1]["content"]
-        for history in model_fn(prefix, history):
+        for history in model_fn(prefix, history, server):
             yield history
-        
-    @spaces.GPU
-    def model_fn(prefix, history):
+
+
+
+
+    def model_fn(prefix, history, server):
+        if server == "Huggingface ZeroGPU":
+            generator = zerogpu_model_fn(prefix, history)
+        elif server == "CPU":
+            generator = cpu_model_fn(prefix, history)
+        # elif server == "RunPod":
+        #     generator = runpod_model_fn(prefix, history)
+        else:
+            raise gr.Error(f"Unknown server: {server}")
+        for history in generator:
+            yield history
+
+    def cpu_model_fn(prefix, history):
         queue = Queue(maxsize=10)
         class MyStreamer:
             def put(self, tokens):
@@ -172,22 +208,44 @@ with gr.Blocks() as demo:
             history[-1]["content"] += text
             yield history
 
-    prefix_box.submit(user_fn, [prefix_box, chatbot_box], [prefix_box, chatbot_box], queue=False).then(
-        generate_fn, chatbot_box, chatbot_box
-    )
+    zerogpu_model_fn = spaces.GPU(cpu_model_fn)
+    
+    def runpod_model_fn(prefix, history):
+        # NOTE
+        runpod_api_key = os.getenv('RUNPOD_API_KEY')
+        runpod_endpoint = os.getenv('RUNPOD_ENDPOINT')
+
+        # synchronized request
+        response = requests.post(
+            f"https://api.runpod.ai/v2/{runpod_endpoint}/runsync",
+            headers={"Authorization": f"Bearer {runpod_api_key}"}, 
+            json={"input": {"prompt": prefix}}
+        ).json()['output'][0]['choices'][0]['tokens'][0]
+        # yield just once
+        history[-1]["content"] += response
+        yield history
+
+
+
+
+
+
+
+
     submit_event = submit_btn.click(user_fn, [prefix_box, chatbot_box], [prefix_box, chatbot_box], queue=False).then(
-        generate_fn, chatbot_box, chatbot_box
+        generate_fn, [chatbot_box, server_box], chatbot_box
     )
     continue_event = continue_btn.click(
-        continue_fn, chatbot_box, chatbot_box
+        continue_fn, [chatbot_box, server_box], chatbot_box
     )
     clear_btn.click(lambda: None, inputs=[], outputs=chatbot_box, cancels=[submit_event, continue_event], queue=False)
 
     def get_audio_fn(history):
         i = random.randint(0, 1000_000_000)
         path = f'./temp/{i}.mid'
+        text = get_last(history)
         try:
-            postprocess(history[-1]["content"], path)
+            postprocess(text, path)
         except Exception as e:
             raise gr.Error(f'Error: {type(e)}, {e}')
         # turn midi into audio with timidity
@@ -201,8 +259,9 @@ with gr.Blocks() as demo:
     def get_midi_fn(history):
         i = random.randint(0, 1000_000_000)
         # turn the text into midi
+        text = get_last(history)
         try:
-            postprocess(history[-1]["content"], f'./temp/{i}.mid')
+            postprocess(text, f'./temp/{i}.mid')
         except Exception as e:
             raise gr.Error(f'Error: {type(e)}, {e}')
         # also render the piano roll