Update app.py

app.py

@@ -10,12 +10,15 @@ 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.checkbox("Text2Image")
 
-if t2i:
+type = st.selectbox("Model Name",("stabilityai/stable-diffusion-2-1-base", "CompVis/stable-diffusion-v1-4"))
+create = st.checkbox("Create The Model")
+
+if create:
     st.title("Text2Image")
-    t2m_mod = create_model()
+    t2m_mod = create_model(loc=type)
     
     prom = st.text_input("# Prompt",'')
 
@@ -34,6 +37,6 @@ if t2i:
 
     create = st.button("Imagine")
     if create:
-        generator = torch.Generator("cuda").manual_seed(int(bu_1))
+        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)