Update app.py

app.py

@@ -32,12 +32,15 @@ def download_prompt_templates():
 def on_token_change(user_token):
     openai.api_key = user_token
 
+def on_type_change(type):
+    print(type)
+
 def on_prompt_template_change(prompt_template):
     if not isinstance(prompt_template, str): return
     return prompt_templates[prompt_template]
 
-def submit_message(user_token, prompt, prompt_template, temperature, max_tokens, context_length, state):
-
+def submit_message(type_select,user_token, prompt, prompt_template, temperature, max_tokens, context_length, state):
+    print(type_select)
     history = state['messages']
 
     if not prompt:
@@ -51,6 +54,14 @@ def submit_message(user_token, prompt, prompt_template, temperature, max_tokens,
 
     prompt_msg = { "role": "user", "content": prompt }
 
+    if not type_select:
+        history.append(prompt_msg)
+        history.append({
+            "role": "system",
+            "content": "Error: Type is not set."
+        })
+        return '', [(history[i]['content'], history[i+1]['content']) for i in range(0, len(history)-1, 2)], f"Total tokens used: 0", state
+
     if not user_token:
         history.append(prompt_msg)
         history.append({
@@ -60,13 +71,25 @@ def submit_message(user_token, prompt, prompt_template, temperature, max_tokens,
         return '', [(history[i]['content'], history[i+1]['content']) for i in range(0, len(history)-1, 2)], f"Total tokens used: 0", state
     
     try:
-        completion = openai.ChatCompletion.create(model="gpt-3.5-turbo", messages=system_prompt + history[-context_length*2:] + [prompt_msg], temperature=temperature, max_tokens=max_tokens)
-
-        history.append(prompt_msg)
-        history.append(completion.choices[0].message.to_dict())
-
-        state['total_tokens'] += completion['usage']['total_tokens']
-    
+        if type_select=='TEXT':
+            completion = openai.ChatCompletion.create(model="gpt-3.5-turbo", messages=system_prompt + history[-context_length*2:] + [prompt_msg], temperature=temperature, max_tokens=max_tokens)
+
+            history.append(prompt_msg)
+            history.append(completion.choices[0].message.to_dict())
+
+            state['total_tokens'] += completion['usage']['total_tokens']
+        elif type_select=='IMAGE':
+            response = openai.Image.create(
+                    prompt=prompt,
+                    n=1,
+                    size="512x512"
+                    )
+            image_url = response['data'][0]['url']
+            history.append(prompt)
+            history.append(image_url)
+
+            # state['total_tokens'] += completion['usage']['total_tokens']
+            
     except Exception as e:
         history.append(prompt_msg)
         history.append({
@@ -105,13 +128,17 @@ with gr.Blocks(css=css) as demo:
                     elem_id="header")
 
         with gr.Row():
-            with gr.Column():
+            with gr.Column(scale=0.7):
                 chatbot = gr.Chatbot(elem_id="chatbox")
-                input_message = gr.Textbox(show_label=False, placeholder="Enter text and press enter", visible=True).style(container=False)
+                with gr.Row():
+                    with gr.Column(scale=0.2, min_width=0):
+                        type_select = gr.Dropdown(show_label=False, choices= ["TEXT", "IMAGE"],value="TEXT",interactive=True)
+                    with gr.Column(scale=0.8):
+                        input_message = gr.Textbox(show_label=False, placeholder="Enter text and press enter", visible=True).style(container=False)
                 btn_submit = gr.Button("Submit")
                 total_tokens_str = gr.Markdown(elem_id="total_tokens_str")
                 btn_clear_conversation = gr.Button("🔃 Start New Conversation")
-            with gr.Column():
+            with gr.Column(scale=0.3):
                 gr.Markdown("Enter your OpenAI API Key. You can get one [here](https://platform.openai.com/account/api-keys).", elem_id="label")
                 user_token = gr.Textbox(value='', placeholder="OpenAI API Key", type="password", show_label=False)
                 prompt_template = gr.Dropdown(label="Set a custom insruction for the chatbot:", choices=list(prompt_templates.keys()))
@@ -124,8 +151,9 @@ with gr.Blocks(css=css) as demo:
     # gr.HTML('''<br><br><br><center>You can duplicate this Space to skip the queue:<a href="https://huggingface.co/spaces/anzorq/chatgpt-demo?duplicate=true"><img src="https://bit.ly/3gLdBN6" alt="Duplicate Space"></a><br>
     #         <p><img src="https://visitor-badge.glitch.me/badge?page_id=anzorq.chatgpt_api_demo_hf" alt="visitors"></p></center>''')
 
-    btn_submit.click(submit_message, [user_token, input_message, prompt_template, temperature, max_tokens, context_length, state], [input_message, chatbot, total_tokens_str, state])
-    input_message.submit(submit_message, [user_token, input_message, prompt_template, temperature, max_tokens, context_length, state], [input_message, chatbot, total_tokens_str, state])
+    type_select.change(on_type_change,inputs=[type_select], outputs=[])
+    btn_submit.click(submit_message, [type_select,user_token, input_message, prompt_template, temperature, max_tokens, context_length, state], [input_message, chatbot, total_tokens_str, state])
+    input_message.submit(submit_message, [type_select,user_token, input_message, prompt_template, temperature, max_tokens, context_length, state], [input_message, chatbot, total_tokens_str, state])
     btn_clear_conversation.click(clear_conversation, [], [input_message, chatbot, total_tokens_str, state])
     prompt_template.change(on_prompt_template_change, inputs=[prompt_template], outputs=[prompt_template_preview])
     user_token.change(on_token_change, inputs=[user_token], outputs=[])
@@ -135,4 +163,6 @@ with gr.Blocks(css=css) as demo:
 
 
 demo.queue(concurrency_count=10)
-demo.launch(auth=("admin", "IBTGeE3NrPsrViDI"),height='800px')
+demo.launch(
+    auth=("admin", "IBTGeE3NrPsrViDI"),
+    height='800px')