Update app.py

app.py

@@ -1,4 +1,4 @@
-from diffusers import StableDiffusionXLPipeline
+from diffusers import StableDiffusionPipeline
 import torch
 from langchain.chains import LLMChain
 from langchain.llms import HuggingFaceHub
@@ -21,6 +21,7 @@ def save_ideas(ideas):
         json.dump(ideas, file)
 
 # Function to generate content
+@torch.no_grad()
 def generate_content(topic):
     hub_llm = HuggingFaceHub(repo_id="HuggingFaceH4/zephyr-7b-beta")
     prompt = PromptTemplate(
@@ -52,51 +53,45 @@ def generate_content(topic):
 
     return content
 
-def make_pipe():
-  pipe = StableDiffusionXLPipeline.from_pretrained("segmind/SSD-1B", torch_dtype=torch.float16, use_safetensors=True, variant="fp16")
-  pipe.to("cuda")
-  return pipe
-
-# 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
-
 
 
-pipe = make_pipe()
-with st.spinner('Please Wait for the models to train...'):
-  if pipe:
-    # Streamlit app
-    st.title("Blog Generator")
+# 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
 
-    # 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('')
+# 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
 # Handle button click
 if button_clicked:
     # Generate content and update existing ideas
-    image = generate_image(pipe,topic)
+    image = generate_image(topic)
     st.image(image)
     content = generate_content(topic)
     existing_ideas.append({topic: content})