Update app.py

app.py

@@ -23,7 +23,8 @@ if 'api_key' not in st.session_state:
     st.session_state['api_key']=""
 
 st.sidebar.title("Settings")
-model_name = st.sidebar.selectbox("Model:", ("gpt-3.5-turbo","gpt-3.5-turbo-16k","gpt-4-32k","falcon-180b-chat","claude-1","claude-2","palm-2","codellama-7b","llama2-7b","llama2-70b","openassistant-llama-30b","mini-orca-v3-13b"))
+model_name = st.sidebar.selectbox("Chat Model:", ("gpt-3.5-turbo","gpt-3.5-turbo-16k","gpt-4","gpt-4-32k","falcon-180b-chat","claude-1","claude-2","palm-2","codellama-7b","llama2-7b","llama2-70b","openassistant-llama-30b","mini-orca-v3-13b"))
+image_model = st.sidebar.selectbox("Image Model:", ("realistic-vision-v1-4","stable-diffusion-v1-4","stable-diffusion-v2-1","kandinsky-2.2","sdxl"))
 counter_placeholder = st.sidebar.empty()
 
 api_key = st.sidebar.text_input("API_Key", value=st.session_state['api_key'], max_chars=None,placeholder=st.session_state['api_key'], key=None, type="password", label_visibility="visible")
@@ -100,7 +101,7 @@ def generate_response(prompt):
     prompt = prompt.lower()
     if prompt.startswith("draw"):
         data = {"prompt": prompt.replace("draw",""),
-                "model": "realistic-vision-v1-4"
+                "model": image_model
                 }
 
         response = requests.post("https://api.webraft.in/api/images/generations", json=data, headers=headers)
@@ -111,7 +112,7 @@ def generate_response(prompt):
 
             response = response.json()
             imageurl = response['data'][0]['url']
-            ie = f"Here is the generated image:\n {imageurl} \n \n Model used: `realistic-vision-v1-4`"
+            ie = f"Here is the generated image:\n {imageurl} \n \n Model used: `{model_name}`"
             st.session_state['messages'].append({"role": "assistant", "content": ie})
 
         except requests.exceptions.JSONDecodeError as e: