from diffusers import DiffusionPipeline
import torch
import streamlit as st

st.title("AI Logo Generator")
pipe = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16, use_safetensors=True, variant="fp16")

prompt = "An astronaut riding a green horse"

images = pipe(prompt=prompt).images[0]

st.image(images, channels="BGR", use_column_width=True)