import streamlit as st
import cv2 as cv
import time
import torch
from diffusers import StableDiffusionPipeline

# "stabilityai/stable-diffusion-2-1-base"
# "CompVis/stable-diffusion-v1-4"
def create_model(loc = "stabilityai/stable-diffusion-2-1-base", mch = 'cpu'):
    pipe = StableDiffusionPipeline.from_pretrained(loc)
    pipe = pipe.to(mch)
    return pipe
    
t2i = st.checkbox("Text2Image")

if t2i:
    st.title("Text2Image")
    t2m_mod = create_model()
    
    prom = st.text_input("# Prompt",'')

    c1,c2,c3 = st.columns([1,1,3])
    c4,c5 = st.columns(2)
    with c1:
      bu_1 = st.text_input("Seed",'999')
    with c2:
      bu_2 = st.text_input("Steps",'12')
    with c3:
      bu_3 = st.text_input("Number of Images",'1')
    with c4:
      sl_1 = st.slider("Width",256,1024,128)
    with c5:
      sl_2 = st.slider("hight",256,1024,128)

    create = st.button("Imagine")
    if create:
        generator = torch.Generator("cuda").manual_seed(int(bu_1))
        img = t2m_mod(prom, width=int(sl_1), height=int(sl_2), num_inference_steps=int(bu_2), generator=generator).images[0]
        st.image(img)