Update app.py

app.py

@@ -18,26 +18,27 @@ the_type = st.selectbox("Model Name",("stabilityai/stable-diffusion-2-1-base",
 create = st.button("Create The Model")
 
 if create:
-    t2m_mod = create_model(loc=the_type)
-    cont = True
-    with cont:
-        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("cpu").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)
+    st.session_state.count.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",256,1024,128)
+with c5:
+  sl_2 = st.slider("hight",256,1024,128)
+
+create = st.button("Imagine")
+if create:
+    generator = torch.Generator("cpu").manual_seed(int(bu_1))
+    model = st.session_state.count.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)