Update app.py
Browse files
app.py
CHANGED
@@ -1,31 +1,43 @@
|
|
1 |
import streamlit as st
|
2 |
-
|
|
|
|
|
|
|
3 |
|
4 |
-
# Define function to generate
|
5 |
-
def generate_image(prompt,
|
6 |
-
|
7 |
-
|
8 |
-
|
9 |
|
10 |
-
#
|
11 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
12 |
|
13 |
return image
|
14 |
|
15 |
# Main function for Streamlit app
|
16 |
def main():
|
17 |
st.title("AI Image Generator")
|
18 |
-
|
19 |
# Input fields
|
20 |
prompt = st.text_input("Enter prompt")
|
21 |
-
|
22 |
-
height = st.number_input("Height", min_value=1, step=1)
|
23 |
|
24 |
if st.button("Generate Image"):
|
25 |
# Check if prompt is provided
|
26 |
if prompt:
|
27 |
# Generate image
|
28 |
-
generated_image = generate_image(prompt,
|
|
|
|
|
29 |
# Display image
|
30 |
st.image(generated_image, caption='Generated Image', use_column_width=True)
|
31 |
else:
|
|
|
1 |
import streamlit as st
|
2 |
+
import torch
|
3 |
+
from diffusers import StableDiffusionXLPipeline, UNet2DConditionModel, EulerDiscreteScheduler
|
4 |
+
from huggingface_hub import hf_hub_download
|
5 |
+
from safetensors.torch import load_file
|
6 |
|
7 |
+
# Define a function to generate the image
|
8 |
+
def generate_image(prompt, num_inference_steps):
|
9 |
+
base = "stable-diffusion"
|
10 |
+
repo = "Geek7/testing"
|
11 |
+
ckpt = "dreamshaper_8.safetensors" # Use the correct ckpt for your step setting!
|
12 |
|
13 |
+
# Load model.
|
14 |
+
unet = UNet2DConditionModel.from_config(base, subfolder="unet").to("cuda", torch.float16)
|
15 |
+
unet.load_state_dict(load_file(hf_hub_download(repo, ckpt), device="cuda"))
|
16 |
+
pipe = StableDiffusionXLPipeline.from_pretrained(base, unet=unet, torch_dtype=torch.float16, variant="fp16").to("cuda")
|
17 |
+
|
18 |
+
# Ensure sampler uses "trailing" timesteps.
|
19 |
+
pipe.scheduler = EulerDiscreteScheduler.from_config(pipe.scheduler.config, timestep_spacing="trailing")
|
20 |
+
|
21 |
+
# Generate image
|
22 |
+
image = pipe(prompt, num_inference_steps=num_inference_steps, guidance_scale=0).images[0]
|
23 |
|
24 |
return image
|
25 |
|
26 |
# Main function for Streamlit app
|
27 |
def main():
|
28 |
st.title("AI Image Generator")
|
29 |
+
|
30 |
# Input fields
|
31 |
prompt = st.text_input("Enter prompt")
|
32 |
+
num_inference_steps = st.slider("Number of Inference Steps", min_value=1, max_value=10, value=2)
|
|
|
33 |
|
34 |
if st.button("Generate Image"):
|
35 |
# Check if prompt is provided
|
36 |
if prompt:
|
37 |
# Generate image
|
38 |
+
generated_image = generate_image(prompt, num_inference_steps)
|
39 |
+
# Save image
|
40 |
+
generated_image.save("output.png")
|
41 |
# Display image
|
42 |
st.image(generated_image, caption='Generated Image', use_column_width=True)
|
43 |
else:
|