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: