# Python
import streamlit as st
import torch
from diffusers import StableCascadeDecoderPipeline, StableCascadePriorPipeline

def load_pipelines(device, dtype):
    prior = StableCascadePriorPipeline.from_pretrained("stabilityai/stable-cascade-prior", torch_dtype=dtype).to(device)
    decoder = StableCascadeDecoderPipeline.from_pretrained("stabilityai/stable-cascade",  torch_dtype=dtype).to(device)
    return prior, decoder

def generate_images(prompt, negative_prompt, num_images_per_prompt, device, dtype):
    with torch.cuda.amp.autocast(dtype=dtype):
        prior_output = prior(
            prompt=prompt,
            height=1024,
            width=1024,
            negative_prompt=negative_prompt,
            guidance_scale=4.0,
            num_images_per_prompt=num_images_per_prompt,
        )
        decoder_output = decoder(
            image_embeddings=prior_output.image_embeddings,
            prompt=prompt,
            negative_prompt=negative_prompt,
            guidance_scale=0.0,
            output_type="pil",
        )
    return decoder_output.images


st.title("Image Generator with Diffusers")

@st.cache(allow_output_mutation=True)
def init_model():
    device = "cuda" if torch.cuda.is_available() else "cpu"
    dtype = torch.bfloat16
    return load_pipelines(device, dtype), device, dtype

(prior, decoder), device, dtype = init_model()

prompt = st.text_input("Enter a prompt:", "Anthropomorphic cat dressed as a pilot")
negative_prompt = st.text_input("Enter a negative prompt:", "")
num_images_per_prompt = st.slider("Number of images per prompt:", 1, 5, 2)


if st.button("Generate"):
    images = generate_images(prompt, negative_prompt, num_images_per_prompt, device, dtype)
    for img in images:
        st.image(img, use_column_width=True)