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) the_type = "stabilityai/stable-diffusion-2-1-base" st.session_state.t2m_mod = create_model(loc=the_type) prom = st.text_input("Prompt",'') neg_prom = st.text_input("Negative Prompt",'') style = st.selectbox("TODO: Image Style",("Cyberpunk", "Picasso", "Real-world specific", "Digital Art", "Aesthetics")) c1,c2,c3,c6 = st.columns([1,1,1,1]) c8 = st.columns([1,1,1,1]) c4,c5 = st.columns(2) with c1: bu_1 = st.text_input("Seed",'666') with c2: bu_2 = st.text_input("Steps",'12') with c3: bu_3 = st.text_input("Number of Images",'1') with c6: bu_6 = st.text_input("Guidance Scale",'7.5') 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, negative_prompt = neg_prom, width=int(sl_1), height=int(sl_2), num_inference_steps=int(bu_2), guidance_scale = float(bu_6), generator=generator).images[0] st.image(IMG) else : PROMS = [prom]*int(bu_3) IMGS = model(PROMS, negative_prompt = neg_prom, width=int(sl_1), height=int(sl_2), num_inference_steps=int(bu_2), guidance_scale = float(bu_6), generator=generator).images st.image(IMGS)