Update app.py

app.py

@@ -175,16 +175,29 @@ def query_summary(text):
     return text
 
 
-def query_generate(text, title, genres, year):
+def query_generate(text, title, genres, year, selected_model='Stable Diffusion v1.5'):
     """
     Get image from HuggingFace Inference API
     -param text: text to generate image
+    -param title: title of the movie
+    -param genres: genres of the movie
+    -param year: year of the movie
+    
     -return: generated image
     """
 
-    #API_URL = "https://api-inference.huggingface.co/models/stabilityai/stable-diffusion-xl-base-1.0"
-    API_URL = "https://api-inference.huggingface.co/models/runwayml/stable-diffusion-v1-5"
-    #API_URL = "https://api-inference.huggingface.co/models/stabilityai/stable-diffusion-2-1"
+    if selected_model=='Stable Diffusion XL':
+        API_URL = "https://api-inference.huggingface.co/models/stabilityai/stable-diffusion-xl-base-1.0"
+
+    elif selected_model=='Stable Diffusion v2.1':
+        API_URL = "https://api-inference.huggingface.co/models/stabilityai/stable-diffusion-2-1"
+
+    elif selected_model=='Stable Diffusion v1.5':
+        API_URL = "https://api-inference.huggingface.co/models/runwayml/stable-diffusion-v1-5"
+        
+    else:
+        raise ValueError("Value not valid for argument 'selected_model'.")
+        
     headers = {"Authorization": f"Bearer {st.secrets['hf_token']}"}
     text = 'A Poster for the movie ' + title.split('(')[0] + 'in portrait mode based on the following synopsis: \"' + text + '\". Style: ' + genres + '. Year ' + year + \
     '. Ignore ' + ''.join(random.choices(string.ascii_letters, k=10))
@@ -205,7 +218,7 @@ def query_generate(text, title, genres, year):
     return response.content
 
 @st.experimental_memo(persist=False, show_spinner=False, suppress_st_warning=True)
-def generate_poster(movie_data):
+def generate_poster(movie_data, selected_model):
     """
     Function for recommending movies
     -param movie_data: metadata of movie selected by user
@@ -240,7 +253,7 @@ def generate_poster(movie_data):
 
     # Get image based on synopsis
     with st.spinner("Generating poster..."):
-        response_content = query_generate(synopsis_sum, title, genres_string, year)
+        response_content = query_generate(synopsis_sum, title, genres_string, year, selected_model)
 
     # Show image
     try: 
@@ -315,8 +328,10 @@ if __name__ == "__main__":
     st.text("")
 
     ## Create button to trigger poster generation
-    buffer1, col1, buffer2 = st.columns([1.3, 1, 1])
-    is_clicked = col1.button(label="Generate poster!")
+    sd_options = ['Stable Diffusion XL', 'Stable Diffusion v2.1', 'Stable Diffusion v1.5']
+    buffer1, col1, col2, buffer2 = st.columns([0.3, 1, 1, 1])
+    session.selected_model = col1.selectbox(label="Select SD model version", options=sd_options)
+    is_clicked = col2.button(label="Generate poster!")
 
     st.text("")
     st.text("")
@@ -327,7 +342,7 @@ if __name__ == "__main__":
 
     ## Generate poster
     if is_clicked:
-        poster = generate_poster(data[data.title_year==session.selected_movie])
+        poster = generate_poster(data[data.title_year==session.selected_movie], session.selected_model)
         generate_poster.clear()
         st.runtime.legacy_caching.clear_cache()