devthedeveloper commited on
Commit
42626a6
·
1 Parent(s): 91d5abc

Create App.py

Browse files
Files changed (1) hide show
  1. app.py +22 -0
app.py ADDED
@@ -0,0 +1,22 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import streamlit as st
2
+ from diffusers import DiffusionPipeline
3
+ import torch
4
+
5
+ # Load the model
6
+ @st.cache(allow_output_mutation=True)
7
+ def load_model():
8
+ pipe = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16, use_safetensors=True, variant="fp16")
9
+ pipe.to("cuda")
10
+ return pipe
11
+
12
+ # Streamlit app
13
+ def main():
14
+ st.title("Image Generation with Diffusion Models")
15
+ prompt = st.text_input("Enter a prompt:", "An astronaut riding a green horse")
16
+ if st.button("Generate"):
17
+ pipe = load_model()
18
+ images = pipe(prompt=prompt).images[0]
19
+ st.image(images, caption=prompt, use_column_width=True)
20
+
21
+ if __name__ == "__main__":
22
+ main()