import streamlit as st

import torch
from torch import autocast
from diffusers import StableDiffusionPipeline


st.set_page_config(layout="wide")

st.title('Play with Stable-Diffusion v1-4')

model_id = "CompVis/stable-diffusion-v1-4"
device = "cuda"

with st.spinner(
	text='Loading...'
):
	pipe = StableDiffusionPipeline.from_pretrained(
		model_id,
		revision="fp16",
		torch_dtype=torch.float16,
		use_auth_token=True
	)

	pipe = pipe.to(device)


def infer(prompt, samples=2, steps=30, scale=7.5, seed=25):
	generator = torch.Generator(device=device).manual_seed(seed)

	with autocast("cuda"):
		images_list = pipe(
			[prompt] * samples,
			num_inference_steps=steps,
			guidance_scale=scale,
			generator=generator,
		)

	images = []
	for image in images_list["sample"]:
		images.append(image)

	return images


with st.form(key='new'):

	prompt = st.text_area(label='Enter prompt')

	col1, col2, col3 = st.columns(3)

	with col1:
		with st.expander(label='Expand parameters'):
			n_samples = st.select_slider(
				label='Num images',
				options=range(1, 5),
				value=1
			)

			steps = st.select_slider(
				label='Steps',
				options=range(1, 61),
				value=40
			)

			scale = st.select_slider(
				label='Guidance Scale',
				options=range(1, 21),
				value=7
			)

	st.form_submit_button()

if prompt:
	st.image(
		infer(
			prompt,
			samples=n_samples,
			steps=steps,
			scale=scale
		),
		caption='result'
	)
else:
	st.warning('Enter prompt.')