Update app.py

app.py

@@ -147,32 +147,28 @@ def pil_image_to_base64(pil_image):
     base64_image = base64.b64encode(buffered.getvalue()).decode('utf-8')
     return base64_image
 
-image_cache = {}
+
 ln0_weight = model.w['blocks.0.ln0.weight'].to(torch.float32).to(device)
 ln0_bias = model.w['blocks.0.ln0.bias'].to(torch.float32).to(device)
-def compute_image_state(image):
-    base64_image = pil_image_to_base64(image)
-    if base64_image in image_cache:
-        image_state = image_cache[base64_image]
-    else:
-        image = image_processor(images=image.convert('RGB'), return_tensors='pt')['pixel_values']
-        image = image.to(device)
-        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=ln0_weight, 
-                                    bias=ln0_bias)
-        _, image_state = model.forward(embs=image_features, state=None)
-        image_cache[base64_image] = image_state
+def compute_image_state(image, prefix_tokens):
+    image = image_processor(images=image.convert('RGB'), return_tensors='pt')['pixel_values']
+    image = image.to(device)
+    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=ln0_weight, 
+                                 bias=ln0_bias)
+    _, image_state = model.forward(tokens=prefix_tokens, embs=image_features, state=None)
     return image_state
 
 def chatbot(image, question):
     if image is None:
         yield "Please upload an image."
         return
-    image_state = compute_image_state(image)
     input_text = generate_prompt(question)
+    prefix_tokens = pipeline.encode(input_text)[-ctx_limit:]
+    image_state = compute_image_state(image, prefix_tokens)
     for output in generate(input_text, image_state):
         yield output