Update app.py

app.py

@@ -17,9 +17,11 @@ the_type = st.selectbox("Model Name",("stabilityai/stable-diffusion-2-1-base",
                                       "CompVis/stable-diffusion-v1-4"))
 create = st.button("Create The Model")
 
+session_state = SessionState.get(name='', mod=None)
+
 if create:
     t2m_mod = create_model(loc=the_type)
-
+    session_state.mod = t2m_mod
 
 prom = st.text_input("# Prompt",'')
 
@@ -39,5 +41,6 @@ with c5:
 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]
+    model = session_state.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)