use inference api

app.py

@@ -2,48 +2,26 @@
 import streamlit as st
 import torch
 from diffusers import StableCascadeDecoderPipeline, StableCascadePriorPipeline
+import requests
+import io
+from PIL import Image
+import os
+from dotenv import load_dotenv
+load_dotenv()
 
-def load_pipelines(device, dtype):
-    prior = StableCascadePriorPipeline.from_pretrained("stabilityai/stable-cascade-prior", torch_dtype=dtype).to(device)
-    decoder = StableCascadeDecoderPipeline.from_pretrained("stabilityai/stable-cascade",  torch_dtype=dtype).to(device)
-    return prior, decoder
-
-def generate_images(prompt, negative_prompt, num_images_per_prompt, device, dtype):
-    with torch.cuda.amp.autocast(dtype=dtype):
-        prior_output = prior(
-            prompt=prompt,
-            height=1024,
-            width=1024,
-            negative_prompt=negative_prompt,
-            guidance_scale=4.0,
-            num_images_per_prompt=num_images_per_prompt,
-        )
-        decoder_output = decoder(
-            image_embeddings=prior_output.image_embeddings,
-            prompt=prompt,
-            negative_prompt=negative_prompt,
-            guidance_scale=0.0,
-            output_type="pil",
-        )
-    return decoder_output.images
-
+API_URL = "https://api-inference.huggingface.co/models/stabilityai/stable-cascade-prior"
+headers = {"Authorization": f"Bearer ${os.getenv('bearer_token')}"}
 
 st.title("Image Generator with Diffusers")
 
-@st.cache(allow_output_mutation=True)
-def init_model():
-    device = "cpu"
-    dtype = torch.bfloat16
-    return load_pipelines(device, dtype), device, dtype
-
-(prior, decoder), device, dtype = init_model()
-
-prompt = st.text_input("Enter a prompt:", "Anthropomorphic cat dressed as a pilot")
-negative_prompt = st.text_input("Enter a negative prompt:", "")
-num_images_per_prompt = st.slider("Number of images per prompt:", 1, 5, 2)
+def query(payload):
+	response = requests.post(API_URL, headers=headers, json=payload)
+	return response.content
 
+prompt = st.text_input("Enter a prompt:", "batman hitting the griddy in gotham")
 
-if st.button("Generate"):
-    images = generate_images(prompt, negative_prompt, num_images_per_prompt, device, dtype)
-    for img in images:
-        st.image(img, use_column_width=True)
+image_bytes = query({
+	"inputs": prompt,
+})
+# You can access the image with PIL.Image for example
+image = Image.open(io.BytesIO(image_bytes))