Shabbir-Anjum commited on
Commit
31b3285
·
verified ·
1 Parent(s): 7cb6427

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +43 -0
app.py ADDED
@@ -0,0 +1,43 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import streamlit as st
2
+ from diffusers import StableDiffusionPipeline
3
+ import torch
4
+ from PIL import Image
5
+ import io
6
+ import base64
7
+
8
+ # Model ID
9
+ model_id = "stabilityai/stable-diffusion-2-1"
10
+
11
+ # Load the model from Hugging Face Hub
12
+ pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
13
+ pipe = pipe.to("cuda" if torch.cuda.is_available() else "cpu")
14
+
15
+ # Streamlit app title and description
16
+ st.title("Stable Diffusion Image Generation")
17
+ st.write("Generate images using the Stability AI's Stable Diffusion model.")
18
+
19
+ # Input for the prompt
20
+ prompt = st.text_input("Enter a prompt for the image:")
21
+
22
+ # Generate button
23
+ if st.button("Generate Image"):
24
+ if prompt:
25
+ with st.spinner('Generating image...'):
26
+ # Generate image
27
+ image = pipe(prompt).images[0]
28
+
29
+ # Convert image to displayable format
30
+ buffered = io.BytesIO()
31
+ image.save(buffered, format="PNG")
32
+ img_str = base64.b64encode(buffered.getvalue()).decode()
33
+
34
+ # Display the image
35
+ st.image(image, caption="Generated Image", use_column_width=True)
36
+
37
+ # Option to download the image
38
+ st.markdown(
39
+ f'<a href="data:image/png;base64,{img_str}" download="generated_image.png">Download Image</a>',
40
+ unsafe_allow_html=True,
41
+ )
42
+ else:
43
+ st.error("Please enter a prompt to generate an image.")