Update app.py

app.py

@@ -1,5 +1,7 @@
 import gradio as gr
 import os, gc
+import base64
+from io import BytesIO
 import torch
 import torch.nn.functional as F
 from transformers import CLIPImageProcessor
@@ -103,17 +105,35 @@ examples = [
     ]
 ]
 
+
+def pil_image_to_base64(pil_image):
+    buffered = BytesIO()
+    pil_image.save(buffered, format="JPEG")  # You can change the format as needed (JPEG, PNG, etc.)
+    # Encodes the image data into base64 format as a bytes object
+    base64_image = base64.b64encode(buffered.getvalue()).decode('utf-8')
+    return base64_image
+
+image_cache = {}
+def get_image_features(image):
+    base64_image = pil_image_to_base64(image)
+    if base64_image in image_cache:
+        image_features = image_cache[base64_image]
+    else:
+        image = image_processor(images=image.convert('RGB'), return_tensors='pt')['pixel_values']
+        image_features = visual_encoder.encode_images(image.unsqueeze(0)).squeeze(0) # [L, D]
+        # apply layer norm to image feature, very important
+        image_features = F.layer_norm(image_features, 
+                                    (image_features.shape[-1],), 
+                                    weight=model.w['blocks.0.ln0.weight'], 
+                                    bias=model.w['blocks.0.ln0.bias'])
+        image_cache[base64_image] = image_features
+    return image_features
+
 def chatbot(image, question):
     if image is None:
         yield "Please upload an image."
         return
-    image = image_processor(images=image.convert('RGB'), return_tensors='pt')['pixel_values']
-    image_features = visual_encoder.encode_images(image.unsqueeze(0)).squeeze(0) # [L, D]
-    # apply layer norm to image feature, very important
-    image_features = F.layer_norm(image_features, 
-                                  (image_features.shape[-1],), 
-                                  weight=model.w['blocks.0.ln0.weight'], 
-                                  bias=model.w['blocks.0.ln0.bias'])
+    image_features = get_image_features(image)
     input_text = generate_prompt(question)
     for output in generate(input_text, image_features):
         yield output