Update app.py

app.py

@@ -4,8 +4,7 @@ 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)
@@ -19,16 +18,18 @@ Txt2Img
 
 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.generator = torch.Generator("cpu").manual_seed(int(bu_1))
 
 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:
@@ -41,26 +42,22 @@ with c5:
   sl_2 = st.slider("hight",128,1024,512,8)
 
 create = st.button("Imagine")
+
 if create:
-    # generator = torch.Generator("cpu").manual_seed(int(bu_1))
-    # model = st.session_state.t2m_mod
-    # img = model(prom, width=int(sl_1), height=int(sl_2), num_inference_steps=int(bu_2), generator=generator).images[0]
-    # st.image(img)
+    model = st.session_state.t2m_mod
+    generator = st.session_state.generator
 
     if int(bu_3) == 1 :
-      generator = torch.Generator("cpu").manual_seed(int(bu_1))
-      model = st.session_state.t2m_mod
       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 :
-      generator = torch.Generator("cpu").manual_seed(int(bu_1))
       PROMS = [prom]*int(bu_3)
-      model = st.session_state.t2m_mod
+        
       IMGS = model(PROMS, width=int(sl_1), height=int(sl_2),
                      num_inference_steps=int(bu_2),
                      generator=generator).images
     
-      # IMGS = np.hstack(IMGS)
       st.image(IMGS)