Spaces:
Running
Running
import streamlit as st | |
import torch | |
from diffusers import StableDiffusionXLPipeline | |
# Load the stable diffusion model | |
pipe = StableDiffusionXLPipeline.from_pretrained( | |
"stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16 | |
) | |
pipe = pipe.to("cuda") # Move the model to GPU for faster processing | |
# Create a Streamlit app | |
st.title("Image Generation App") | |
# Add a text input for the prompt | |
prompt = st.text_input("Prompt", "a photo of an astronaut riding a horse on mars") | |
# Generate the image based on the prompt | |
if st.button("Generate Image"): | |
# Use the model to generate the image | |
image = pipe(prompt).images[0] | |
# Display the generated image | |
st.image(image, caption="Generated Image", use_column_width=True) | |