Update app.py

app.py

@@ -14,6 +14,8 @@ claude_models = ["claude-3-opus-20240229", "claude-3-sonnet-20240229", "claude-3
 # List of available OpenAI models
 openai_models = ["gpt-4", "gpt-4-32k", "gpt-3.5-turbo", "gpt-4-0125-preview", "gpt-4-turbo-preview", "gpt-4-1106-preview", "gpt-4-0613"]
 
+image_prompter = ["SDXL", "midjourney"]
+
 both_models = claude_models + openai_models
 
 def generate_response(endpoint, api_key, model, user_prompt):
@@ -113,15 +115,15 @@ def extract_json(generated_output):
         print(e)
         return None, None
 
-def generate_second_response(endpoint, api_key, model, generated_output):
+def generate_second_response(endpoint, api_key, model, generated_output, image_model):
     if endpoint in default_urls:
         #check api keys as normal
         if api_key.startswith("sk-ant-"):
             client = Anthropic(api_key=api_key, base_url=endpoint)
-            system_prompt_path = __file__.replace("app.py", "diffusion.txt")
+            system_prompt_path = __file__.replace("app.py", f"{image_model}.txt")
         elif api_key.startswith("sk-"):
             client = OpenAI(api_key=api_key, base_url=endpoint)
-            system_prompt_path = __file__.replace("app.py", "diffusion.txt")
+            system_prompt_path = __file__.replace("app.py", f"{image_model}.txt")
         else:
             print("Invalid API key")
             return "Invalid API key", "Invalid API key", None
@@ -129,11 +131,11 @@ def generate_second_response(endpoint, api_key, model, generated_output):
         if model in claude_models:
             # Set the Anthropic API key
             client = Anthropic(api_key=api_key, base_url=endpoint)
-            system_prompt_path = __file__.replace("app.py", "diffusion.txt")
+            system_prompt_path = __file__.replace("app.py", f"{image_model}.txt")
         else:
             # Set the OpenAI API key
             client = OpenAI(api_key=api_key, base_url=endpoint)
-            system_prompt_path = __file__.replace("app.py", "diffusion.txt")
+            system_prompt_path = __file__.replace("app.py", f"{image_model}.txt")
 
     # Read the system prompt from a text file
     with open(system_prompt_path, "r") as file:
@@ -192,6 +194,7 @@ with gr.Blocks() as demo:
 
     with gr.Row():
         with gr.Column():
+            image_model = gr.Dropdown(choices=image_prompter, label="Image Model to prompt for", value="SDXL")
             generate_button_2 = gr.Button("Generate SDXL Prompt")
 
         with gr.Column():
@@ -208,6 +211,6 @@ with gr.Blocks() as demo:
     api_key.change(update_models, inputs=api_key, outputs=[model_dropdown, endpoint])
 
     generate_button.click(generate_response, inputs=[endpoint, api_key, model_dropdown, user_prompt], outputs=[generated_output, json_output, json_download])
-    generate_button_2.click(generate_second_response, inputs=[endpoint, api_key, model_dropdown, generated_output], outputs=generated_output_2)
+    generate_button_2.click(generate_second_response, inputs=[endpoint, api_key, model_dropdown, generated_output, image_model], outputs=generated_output_2)
 
 demo.launch()