Add option to select GPU

app.py

@@ -21,7 +21,12 @@ with (gr.Blocks() as app):
                 with gr.Row():
                     output_audio = gr.Audio(interactive=False, label='Audio', autoplay=True)
                 with gr.Row():
-                    checkbox = gr.Checkbox(label='Get video', info='Remember that this has a cost')
+                    checkbox = gr.Checkbox(
+                        label='Get video', info='Remember that this has a cost'
+                    )
+                    radio = gr.Radio(
+                        choices=['small', 'big'], value='small', label='GPU', info='Select the size of GPU'
+                    )
 
             with gr.Column():
                 with gr.Row():
@@ -59,13 +64,13 @@ with (gr.Blocks() as app):
 
     button_text.click(
         get_answer_text,
-        [text, chat, user_id, checkbox],
+        [text, chat, user_id, checkbox, radio],
         [video, output_audio, chat, text]
     )
 
     button_audio.click(
         get_answer_audio,
-        [audio, chat, user_id, checkbox],
+        [audio, chat, user_id, checkbox, radio],
         [video, output_audio, chat, audio]
     )
 

functions.py

@@ -7,14 +7,6 @@ import gradio as gr
 from datetime import datetime
 from huggingface_hub import hf_hub_download, HfApi
 
-API_TOKEN = os.getenv('API_TOKEN')
-API_URL = os.getenv('API_URL')
-
-headers = {
-    "Authorization": f"Bearer {API_TOKEN}",
-    "Content-Type": "application/json"
-}
-
 
 def get_main_data():
     """
@@ -43,10 +35,18 @@ def make_visible():
     return gr.Row.update(visible=True), gr.Row.update(visible=True)
 
 
-def _query(payload):
+def _query(payload, size_gpu: str):
     """
     Returns the json from a post request. It is done to the BellaAPI
     """
+    API_TOKEN = os.getenv(f'API_TOKEN')
+    API_URL = os.getenv(f'API_URL_{size_gpu}')
+
+    headers = {
+        "Authorization": f"Bearer {API_TOKEN}",
+        "Content-Type": "application/json"
+    }
+
     response = requests.post(API_URL, headers=headers, json=payload)
     return response.json()
 
@@ -71,7 +71,7 @@ def init_chatbot(chatbot: list[tuple[str, str]]):
         "date": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
         'get_video': True
     }}
-    output = _query(inputs)
+    output = _query(inputs, 'small')
 
     chatbot.append(('', output['answer']))
     _download_media(output['link_media'], 'video')
@@ -79,7 +79,9 @@ def init_chatbot(chatbot: list[tuple[str, str]]):
     return 'video.mp4', chatbot, output['user_id']
 
 
-def get_answer_text(question: str, chatbot: list[tuple[str, str]], user_id: str, checkbox: bool):
+def get_answer_text(
+        question: str, chatbot: list[tuple[str, str]], user_id: str, checkbox: bool, size_gpu: str
+):
     """
     Gets the answer of the chatbot
     """
@@ -87,11 +89,13 @@ def get_answer_text(question: str, chatbot: list[tuple[str, str]], user_id: str,
     inputs = {'inputs': {
         'text': question, 'user_id': user_id, 'get_video': checkbox
     }}
-    output = _query(inputs)
+    output = _query(inputs, size_gpu)
     return _update_elements(question, chatbot, output, checkbox, '')
 
 
-def get_answer_audio(audio_path, chatbot: list[tuple[str, str]], user_id: str, checkbox: bool):
+def get_answer_audio(
+        audio_path, chatbot: list[tuple[str, str]], user_id: str, checkbox: bool, size_gpu: str
+):
     """
     Gets the answer of the chatbot
     """
@@ -104,7 +108,7 @@ def get_answer_audio(audio_path, chatbot: list[tuple[str, str]], user_id: str, c
     inputs = {'inputs': {
         'is_audio': True, 'audio': encoded_audio, 'user_id': user_id, 'get_video': checkbox
     }}
-    output = _query(inputs)
+    output = _query(inputs, size_gpu)
 
     # Transcription of the audio
     question = output['question']