Update app.py

app.py

@@ -3,6 +3,7 @@ import torch
 from langchain.chains import LLMChain
 from langchain.llms import HuggingFaceHub
 from langchain.prompts import PromptTemplate
+import requests
 import streamlit as st
 import json
 
@@ -20,13 +21,9 @@ def save_ideas(ideas):
     with open("ideas.json", "w") as file:
         json.dump(ideas, file)
 
-# 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)
+# Function to generate content
+@torch.no_grad()
+def generate_content(topic):
     hub_llm = HuggingFaceHub(repo_id="HuggingFaceH4/zephyr-7b-beta")
     prompt = PromptTemplate(
         input_variables=['keyword'],
@@ -37,16 +34,6 @@ def load_models():
         """
     )
     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 = [
@@ -70,12 +57,25 @@ def generate_content(topic):
 
 
 # generate image
+
+import io
+from PIL import Image
+# Function to generate an image using the pre-created or newly created pipeline
 @torch.no_grad()
 def generate_image(topic):
-    prompt = f"blog banner about {topic}"
-    image = pipe(prompt).images[0]
-    return image
+  API_URL = "https://api-inference.huggingface.co/models/runwayml/stable-diffusion-v1-5"
+  headers = {"Authorization": "Bearer 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"}
+
+  def query(payload):
+    response = requests.post(API_URL, headers=headers, json=payload)
+    return response.content
+  image_bytes = query({
+    "inputs": f"A blog banner about {topic}",
+  })
+  # You can access the image with PIL.Image for example
 
+  image = Image.open(io.BytesIO(image_bytes))
+  return image
 
 # Streamlit app
 st.title("Blog Generator")
@@ -102,11 +102,13 @@ else:
 # Handle button click
 if button_clicked:
     # Generate content and update existing ideas
-    content, image_path = generate_content(topic),generate_image(topic)
+    content, image = generate_content(topic),generate_image(topic)
+    if image:
+      image_path = f"{topic}.png"
     existing_ideas.append({topic: {"content": content, "image_path": image_path}})
     save_ideas(existing_ideas)
     # 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=f"selectbox{topic}", index=keys.index(topic))
-    st.markdown(content)
-    st.image(image_path)
+    st.image(image)
+    st.markdown(content)