| 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) | |
| 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) |