import streamlit as st
import torch
from diffusers import StableDiffusion3Pipeline
import random

# Load model
device = "cuda" if torch.cuda.is_available() else "cpu"
pipe = StableDiffusion3Pipeline.from_pretrained(
    "stabilityai/stable-diffusion-3-medium-diffusers",
    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
).to(device)

st.title("🎨 Stable Diffusion 3 Medium - Streamlit")

prompt = st.text_input("Prompt", "Astronaut in a jungle, cold color palette, 8k")
negative_prompt = st.text_input("Negative Prompt", "blurry, low quality, text")
guidance_scale = st.slider("Guidance Scale", 0.0, 10.0, 5.0)
num_inference_steps = st.slider("Inference Steps", 1, 50, 28)
width = st.slider("Image Width", 256, 1344, 1024, step=64)
height = st.slider("Image Height", 256, 1344, 1024, step=64)
random_seed = st.checkbox("Randomize Seed", value=True)
seed = random.randint(0, 2**32 - 1) if random_seed else st.number_input("Seed", value=0, step=1)

if st.button("Generate Image"):
    with st.spinner("Generating..."):
        generator = torch.Generator().manual_seed(seed)
        image = pipe(
            prompt=prompt,
            negative_prompt=negative_prompt,
            guidance_scale=guidance_scale,
            num_inference_steps=num_inference_steps,
            width=width,
            height=height,
            generator=generator
        ).images[0]
        st.image(image, caption="Generated Image")
        st.write(f"Seed used: {seed}")