import streamlit as st
from diffusers import StableDiffusionPipeline
import torch

# Load the model
repo = "IDKiro/sdxs-512-0.9"
seed = 42
pipe = StableDiffusionPipeline.from_pretrained(repo, torch_dtype=torch.float32)

# Streamlit app
st.title("Image generation demo")

prompt = st.text_input("Enter a text prompt:", "A gold ring")

if st.button("Generate Image"):
    # Generate image from the prompt
    image = pipe(
        prompt, 
        num_inference_steps=1, 
        guidance_scale=0,
        generator=torch.manual_seed(seed)
    ).images[0]
    
    # Display the image
    st.image(image, caption="Generated Image", channels="RGB")