Update app.py

app.py

@@ -5,16 +5,15 @@ import torch
 import numpy as np
 import urllib.request
 
-# Load the model without caching to avoid serialization issues
+memory = {}
+
+@st.cache_resource
 def load_model():
     model = AutoModel.from_pretrained("ragavsachdeva/magi", trust_remote_code=True)
     device = "cuda" if torch.cuda.is_available() else "cpu"
     model.to(device)
     return model
 
-# Initialize the model once at the top level, outside any caching functions
-model = load_model()
-
 @st.cache_data
 def read_image_as_np_array(image_path):
     if "http" in image_path:
@@ -70,6 +69,12 @@ def predict_ocr(
         ocr_results = model.predict_ocr([image], text_bboxes_for_all_images)
     return ocr_results
 
+model = load_model()
+
+# Add a button to clear memory
+if st.button("Clear Memory"):
+    memory.clear()
+
 # Streamlit UI elements
 st.markdown("""    
 <style> .title-container { background-color: #0d1117; padding: 20px; border-radius: 10px; margin: 20px; }