Update app.py

app.py

@@ -35,11 +35,11 @@ image_processor = CLIPImageProcessor.from_pretrained(vision_tower_name)
 ##########################################################################
 def generate_prompt(instruction):
     instruction = instruction.strip().replace('\r\n','\n').replace('\n\n','\n')
-    return f"{instruction}\n\nAssistant:"
+    return f"\n{instruction}\n\nAssistant:"
 
 def generate(
     ctx,
-    image_features,
+    image_state,
     token_count=128,
     temperature=0.2,
     top_p=0.3,
@@ -58,10 +58,8 @@ def generate(
     occurrence = {}
     for i in range(int(token_count)):
         if i == 0:
-            input_ids = pipeline.encode(ctx)
-            text_embs = model.w['emb.weight'][input_ids]
-            input_embs = torch.cat((image_features, text_embs), dim=0)[-ctx_limit:]
-            out, state = model.forward(embs=input_embs, state=None)
+            input_ids = pipeline.encode(ctx)[-ctx_limit:]
+            out, state = model.forward(tokens=input_ids, state=image_state)
         else:
             input_ids = [token]
             out, state = model.forward(tokens=input_ids, state=state)
@@ -113,11 +111,10 @@ def pil_image_to_base64(pil_image):
     return base64_image
 
 image_cache = {}
-def get_image_features(image):
+def compute_image_state(image):
     base64_image = pil_image_to_base64(image)
     if base64_image in image_cache:
-        image_features = image_cache[base64_image]
-        print(f"use cache {base64_image[:10]}")
+        image_state = 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]
@@ -126,16 +123,17 @@ def get_image_features(image):
                                     (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
+        _, image_state = model.forward(embs=image_features, state=None)
+        image_cache[base64_image] = image_state
+    return image_state
 
 def chatbot(image, question):
     if image is None:
         yield "Please upload an image."
         return
-    image_features = get_image_features(image)
+    image_state = compute_image_state(image)
     input_text = generate_prompt(question)
-    for output in generate(input_text, image_features):
+    for output in generate(input_text, image_state):
         yield output
 
 with gr.Blocks(title=title) as demo: