move back to streamlit 1.2.0

image2text.py

@@ -3,6 +3,7 @@ from text2image import get_model, get_tokenizer, get_image_transform
 from utils import text_encoder, image_encoder
 from PIL import Image
 from jax import numpy as jnp
+from io import BytesIO
 import pandas as pd
 import requests
 import jax
@@ -29,7 +30,7 @@ def app():
 
     image_url = st.text_input(
         "You can input the URL of an image",
-        value="https://upload.wikimedia.org/wikipedia/commons/b/bc/Juvenile_Ragdoll.jpg",
+        value="https://upload.wikimedia.org/wikipedia/commons/thumb/8/88/Ragdoll%2C_blue_mitted.JPG/1280px-Ragdoll%2C_blue_mitted.JPG",
     )
 
     MAX_CAP = 4
@@ -59,17 +60,13 @@ def app():
 
                 text_embeds = list()
                 for i, c in enumerate(captions):
-                    text_embeds.extend(text_encoder(c, model, tokenizer))
+                    text_embeds.extend(text_encoder(c, model, tokenizer)[0])
 
                 text_embeds = jnp.array(text_embeds)
-                image_raw = requests.get(
-                    image_url,
-                    stream=True,
-                ).raw
-
-                image = Image.open(image_raw).convert("RGB")
+                response = requests.get(image_url)
+                image = Image.open(BytesIO(response.content)).convert("RGB")
                 transform = get_image_transform(model.config.vision_config.image_size)
-                image_embed = image_encoder(transform(image), model)
+                image_embed, _ = image_encoder(transform(image), model)
 
                 # we could have a softmax here
                 cos_similarities = jax.nn.softmax(
@@ -87,9 +84,6 @@ def app():
         gc.collect()
 
     elif image_url:
-        image_raw = requests.get(
-            image_url,
-            stream=True,
-        ).raw
-        image = Image.open(image_raw).convert("RGB")
+        response = requests.get(image_url)
+        image = Image.open(BytesIO(response.content)).convert("RGB")
         st.image(image)

utils.py

@@ -48,7 +48,7 @@ def image_encoder(image, model):
     features = model.get_image_features(image,)
     norms = jnp.linalg.norm(features, axis=-1, keepdims=True)
     features = features / norms
-    return features
+    return features, norms
 
 
 def precompute_image_features(model, loader):
@@ -62,8 +62,7 @@ def precompute_image_features(model, loader):
 
 
 def find_image(text_query, model, dataset, tokenizer, image_features, n, dataset_name):
-    zeroshot_weights = text_encoder(text_query, model, tokenizer)
-    zeroshot_weights /= jnp.linalg.norm(zeroshot_weights)
+    zeroshot_weights, _ = text_encoder(text_query, model, tokenizer)
     distances = jnp.dot(image_features, zeroshot_weights.reshape(-1, 1))
     file_paths = []
     for i in range(1, n + 1):