models / app.py
Geek7's picture
Upload app.py
c71f7b7
raw
history blame
796 Bytes
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("cuda")
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.")