Add gc.collect

image2text.py

@@ -6,6 +6,8 @@ from jax import numpy as jnp
 import pandas as pd
 import requests
 import jax
+import gc
+
 
 def app():
     st.title("From Image to Text")
@@ -27,8 +29,8 @@ def app():
 
     image_url = st.text_input(
         "You can input the URL of an image",
-        value="https://www.petdetective.it/wp-content/uploads/2016/04/gatto-toilette.jpg")
-
+        value="https://www.petdetective.it/wp-content/uploads/2016/04/gatto-toilette.jpg",
+    )
 
     MAX_CAP = 4
 
@@ -67,7 +69,9 @@ def app():
                 image_embed = image_encoder(transform(image), model)
 
                 # we could have a softmax here
-                cos_similarities = jax.nn.softmax(jnp.matmul(image_embed, text_embeds.T))
+                cos_similarities = jax.nn.softmax(
+                    jnp.matmul(image_embed, text_embeds.T)
+                )
 
                 chart_data = pd.Series(cos_similarities[0], index=captions)
 
@@ -77,3 +81,5 @@ def app():
 
                 with col2:
                     st.image(image)
+
+        gc.collect()

text2image.py

@@ -3,6 +3,7 @@ import os
 import requests
 import zipfile
 import natsort
+import gc
 
 os.environ["TOKENIZERS_PARALLELISM"] = "false"
 from stqdm import stqdm
@@ -144,13 +145,11 @@ def app():
     with col2:
         dataset_name = st.selectbox("IR dataset", ["Unsplash", "CC"])
 
-    query = query if query else suggestions[sugg_idx] if sugg_idx > -1 else ""
+    query = suggestions[sugg_idx] if sugg_idx > -1 else query if query else ""
 
     if query:
         with st.spinner("Computing..."):
 
-            model = get_model()
-
             if dataset_name == "Unsplash":
                 download_images()
 
@@ -173,3 +172,7 @@ def app():
             )
 
         st.image(image_paths)
+
+        gc.collect()
+
+        sugg_idx = -1