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


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.title("""
# Txt2Img
# ###### `CLICK "Create_Update_Model"` :
# - `FIRST RUN OF THE CODE`
# - `CHANGING MODEL`""")

# the_type = st.selectbox("Model",("stabilityai/stable-diffusion-2-1-base",
#                                       "CompVis/stable-diffusion-v1-4"))

# create = st.button("Create The Model")

# if create:
#     st.session_state.t2m_mod = create_model(loc=the_type)

st.session_state.t2m_mod = create_model(loc=the_type)


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",128,1024,512,8)
with c5:
  sl_2 = st.slider("hight",128,1024,512,8)

st.session_state.generator = torch.Generator("cpu").manual_seed(int(bu_1))

create = st.button("Imagine")

if create:
    model = st.session_state.t2m_mod
    generator = st.session_state.generator

    if int(bu_3) == 1 :
      IMG = model(prom, width=int(sl_1), height=int(sl_2),
                    num_inference_steps=int(bu_2),
                    generator=generator).images[0]
      st.image(IMG)
        
    else :
      PROMS = [prom]*int(bu_3)
        
      IMGS = model(PROMS, width=int(sl_1), height=int(sl_2),
                     num_inference_steps=int(bu_2),
                     generator=generator).images
    
      st.image(IMGS)