fix: Load model upfront

app.py

@@ -9,6 +9,20 @@ from model import load_clip_model, load_gpt_model, load_model
 # Configure Streamlit page
 st.set_page_config(page_title="Caption Machine", page_icon="💥")
 
+# Set Session
+
+model, image_transform, tokenizer = load_model()
+
+if 'model' not in st.session_state:
+    st.session_state['model'] = model
+
+if 'image_transform' not in st.session_state:
+    st.session_state['image_transform'] = image_transform
+
+if 'tokenizer' not in st.session_state:
+    st.session_state['tokenizer'] = tokenizer
+
+
 
 # Force responsive layout for columns also on mobile
 st.write(
@@ -44,9 +58,7 @@ if upload_file is not None:
     st.write("Image Uploaded Successfully")
 
     # gpt_model, tokenizer = load_gpt_model()
-
-    model, image_transform, tokenizer = load_model()
-    caption = generate_text(model, img, tokenizer, image_transform)
+    caption = generate_text(st.session_state['model'], img, st.session_state['tokenizer'], st.session_state['image_transform'])
 
     st.write(caption)