Update app.py

app.py

@@ -1,21 +1,27 @@
 import streamlit as st
-from diffusers import DiffusionPipeline
-import torch
+import requests
 from PIL import Image
-from io import BytesIO
+import io
+import os
 
-# Load the diffusion pipeline
-@st.cache_resource
-def load_pipeline():
-    return DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0").to("cuda" if torch.cuda.is_available() else "cpu")
+# Set up your Hugging Face API key
+token = os.getenv("HF_TOKEN")
 
-pipeline = load_pipeline()
+API_URL = "https://api-inference.huggingface.co/models/stabilityai/stable-diffusion-xl-base-1.0"
+headers = {"Authorization": f"Bearer {token}"}
+
+def query(payload):
+    response = requests.post(API_URL, headers=headers, json=payload)
+    if response.status_code != 200:
+        st.error(f"Error: {response.status_code} - {response.text}")
+        return None
+    return response.content
 
 def generate_image(prompt):
-    # Generate image
-    with torch.no_grad():
-        result = pipeline(prompt).images[0]
-    return result
+    image_bytes = query({"inputs": prompt})
+    if image_bytes:
+        return Image.open(io.BytesIO(image_bytes))
+    return None
 
 def main():
     st.title("Stable Diffusion Image Generator")
@@ -27,7 +33,8 @@ def main():
         if prompt:
             # Generate and display the image
             image = generate_image(prompt)
-            st.image(image, caption="Generated Image")
+            if image:
+                st.image(image, caption="Generated Image")
         else:
             st.warning("Please enter a prompt.")