File size: 877 Bytes
c71f7b7 a96537b c71f7b7 e0756bc c71f7b7 e0756bc |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 |
import streamlit as st
from diffusers import AutoPipelineForText2Image
import torch
# Load model and move it to CUDA
pipe = AutoPipelineForText2Image.from_pretrained("stabilityai/sdxl-turbo", torch_dtype=torch.float16, variant="fp16")
pipe.to("cpu")
# Cache the image generation function
@st.cache(suppress_st_warning=True)
def generate_image(prompt):
# Generate image using the model
image = pipe(prompt=prompt, num_inference_steps=1, guidance_scale=0.0).images[0]
return image
# Streamlit app
st.title("Text to Image Generation App")
# User input prompt
prompt = st.text_area("Enter a prompt for image generation:")
if st.button("Generate Image"):
if prompt:
# Generate and display the image
st.image(generate_image(prompt), caption="Generated Image", use_column_width=True)
else:
st.warning("Please enter a prompt.")
|