Update app.py

app.py

@@ -20,9 +20,13 @@ def save_ideas(ideas):
     with open("ideas.json", "w") as file:
         json.dump(ideas, file)
 
-# Function to generate content
-@torch.no_grad()
-def generate_content(topic):
+# Create the pipeline and langchain model 
+def load_models():
+  with torch.no_grad():
+    model_id = "CompVis/stable-diffusion-v1-4"
+    device = "cuda"
+    pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
+    pipe = pipe.to(device)
     hub_llm = HuggingFaceHub(repo_id="HuggingFaceH4/zephyr-7b-beta")
     prompt = PromptTemplate(
         input_variables=['keyword'],
@@ -33,6 +37,16 @@ def generate_content(topic):
         """
     )
     hub_chain = LLMChain(prompt=prompt, llm=hub_llm, verbose=True)
+    return hub_chain,pipe
+
+# Wait for the models to be created
+with st.spinner("Creating generation pipelines. Please Wait:)..."):
+  hub_chain,pipe = load_models()
+
+
+# Function to generate content
+@torch.no_grad()
+def generate_content(topic):
     content = hub_chain.run(topic)
 
     subheadings = [
@@ -58,10 +72,6 @@ def generate_content(topic):
 # generate image
 @torch.no_grad()
 def generate_image(topic):
-    model_id = "CompVis/stable-diffusion-v1-4"
-    device = "cuda"
-    pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
-    pipe = pipe.to(device)
     prompt = f"blog banner about {topic}"
     image = pipe(prompt).images[0]
     return image
@@ -82,22 +92,21 @@ st.sidebar.header("Previous Ideas:")
 keys = list(set([key for idea in existing_ideas for key in idea.keys()]))
 if topic in keys:
     index = keys.index(topic)
-    selected_idea = st.sidebar.selectbox("Select Idea", keys, key="selectbox", index=index)
-    # Display content for the selected idea
+    selected_idea = st.sidebar.selectbox("Select Idea", keys, key=f"selectbox{topic}", index=index)
+    # Display content and image for the selected idea
     selected_idea_from_list = next((idea for idea in existing_ideas if selected_idea in idea), None)
-    st.markdown(selected_idea_from_list[selected_idea])
+    st.markdown(selected_idea_from_list[selected_idea]["content"])
+    st.image(selected_idea_from_list[selected_idea]["image_path"])
 else:
     index = 0
 # Handle button click
 if button_clicked:
     # Generate content and update existing ideas
-    image = generate_image(topic)
-    st.image(image)
-    content = generate_content(topic)
-    existing_ideas.append({topic: content})
+    content, image_path = generate_content(topic),generate_image(topic)
+    existing_ideas.append({topic: {"content": content, "image_path": image_path}})
     save_ideas(existing_ideas)
-    st.experimental_rerun()
     # Update keys and selected idea in the sidebar
     keys = list(set([key for idea in existing_ideas for key in idea.keys()]))
-    selected_idea = st.sidebar.selectbox("Select Idea", keys, key="selectbox", index=keys.index(topic))
+    selected_idea = st.sidebar.selectbox("Select Idea", keys, key=f"selectbox{topic}", index=keys.index(topic))
     st.markdown(content)
+    st.image(image_path)