Update app.py

app.py

@@ -1,3 +1,5 @@
+from diffusers import StableDiffusionXLPipeline
+import torch
 from langchain.chains import LLMChain
 from langchain.llms import HuggingFaceHub
 from langchain.prompts import PromptTemplate
@@ -50,36 +52,57 @@ def generate_content(topic):
 
     return content
 
-# Streamlit app
-st.title("Blog Generator")
+def make_pipe():
+  pipe = StableDiffusionXLPipeline.from_pretrained("segmind/SSD-1B", torch_dtype=torch.float16, use_safetensors=True, variant="fp16")
+  pipe.to("cuda")
+  return pipe
 
-# Input and button
-topic = st.text_input("Enter Title for the blog")
-button_clicked = st.button("Create blog!")
+# generate image
+def generate_image(pipe,topic):
+  prompt = f"A banner for a blog about{topic}" # Your prompt here
+  neg_prompt = "ugly, blurry, poor quality" # Negative prompt here
+  image = pipe(prompt=prompt, negative_prompt=neg_prompt).images[0]
+  return image
 
-# Load existing ideas
-existing_ideas = load_ideas()
-st.sidebar.header("Previous Ideas:")
 
-# Display existing ideas 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=0)
 
-# Display content 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])
+pipe = make_pipe()
+with st.spinner('Please Wait for the models to train...'):
+  if pipe:
+    # Streamlit app
+    st.title("Blog Generator")
+
+    # Input and button
+    topic = st.text_input("Enter Title for the blog")
+    button_clicked = st.button("Create blog!")
+    st.subheader(topic)
+    # Load existing ideas
+    existing_ideas = load_ideas()
+    st.sidebar.header("Previous Ideas:")
+
+    # Display existing ideas in the sidebar
+    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_from_list = next((idea for idea in existing_ideas if selected_idea in idea), None)
+      st.markdown(selected_idea_from_list[selected_idea])
+    else:
+      index = 0
+  st.success('')
+
 
 # Handle button click
 if button_clicked:
     # Generate content and update existing ideas
+    image = generate_image(pipe,topic)
+    st.image(image)
     content = generate_content(topic)
     existing_ideas.append({topic: content})
     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))
-
-    # Update selected idea content
-    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(content)