Update app.py

app.py

@@ -11,14 +11,17 @@ def create_model(loc = "stabilityai/stable-diffusion-2-1-base", mch = 'cpu'):
     pipe = pipe.to(mch)
     return pipe
 
-t2i = st.title("Text2Image")
+t2i = st.title("""# Text2Image
+### CHECK "Create_Update_Model" :
+- FIRST RUN OF THE CODE
+- CHANGING MODEL""")
 
 the_type = st.selectbox("Model Name",("stabilityai/stable-diffusion-2-1-base",
                                       "CompVis/stable-diffusion-v1-4"))
 create = st.button("Create The Model")
 
 if create:
-    st.session_state.count.t2m_mod = create_model(loc=the_type)
+    st.session_state.t2m_mod = create_model(loc=the_type)
 
 
 prom = st.text_input("# Prompt",'')
@@ -39,6 +42,6 @@ with c5:
 create = st.button("Imagine")
 if create:
     generator = torch.Generator("cpu").manual_seed(int(bu_1))
-    model = st.session_state.count.t2m_mod
+    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)