added submit button

app.py

@@ -1,21 +1,46 @@
-# Import convention
 import streamlit as st
+import time
 from diffusers import DiffusionPipeline
 import matplotlib.pyplot as plt
 import torch
 
+if 'button_clicked' not in st.session_state:
+    st.session_state.button_clicked = False
+
+# Define a function to handle the button click
+def on_button_click():
+    st.session_state.button_clicked = True
+
 organ = st.selectbox('Organ', ['Brain', 'Thorax'], index=None)
 modality = st.selectbox('Modality', ['Magnetic Resonance Imaging', 'Computed Tomography'], index=None)
 style = st.selectbox('Style', ['Picasso', 'Van Gogh'], index=None)
 
 prompt_lst = [organ, modality, style]
+
 if None not in prompt_lst:
-    prompt = ','.join(prompt_lst) 
-    print(prompt)
+    st.session_state.button_disabled = False
+else:
+    st.session_state.button_disabled = True
+
+
+if st.session_state.button_clicked:
+    st.session_state.button_disabled = True
+    st.session_state.button_clicked = False
+    st.button('Submit', disabled=st.session_state.button_disabled)
+    with st.spinner('Processing...'):
+        prompt = ','.join(prompt_lst) 
+        print(prompt)
+
+        pipe = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16, variant="fp16")
+        pipe.to("cuda")
 
-    pipe = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16, variant="fp16")
-    pipe.to("cuda")
+        prompt += " high resolution, photorealistic"
+        image = pipe(prompt=prompt).images[0]
+    
+    st.image(image)
+    
+    st.session_state.button_disabled = False
+    
 
-    prompt += " high resolution, photorealistic"
-    image = pipe(prompt=prompt).images[0]
-    st.image(image)
+else:
+    st.button('Submit', on_click=on_button_click, disabled=st.session_state.button_disabled)