Add prefilling

app.py

@@ -85,35 +85,51 @@ def deepseek(image, text_input, model_id):
         system_prompt=""
     ).to(vl_gpt.device)
 
-    # run image encoder to get the image embeddings
-    inputs_embeds = vl_gpt.prepare_inputs_embeds(**prepare_inputs)
-
-    # run the model to get the response
-    outputs = vl_gpt.language.generate(
-        inputs_embeds=inputs_embeds,
-        attention_mask=prepare_inputs.attention_mask,
-        pad_token_id=tokenizer.eos_token_id,
-        bos_token_id=tokenizer.bos_token_id,
-        eos_token_id=tokenizer.eos_token_id,
-        max_new_tokens=512,
-        do_sample=False,
-        use_cache=True
-    )
-
-    answer = tokenizer.decode(outputs[0].cpu().tolist(), skip_special_tokens=False)
-
-    print(f"{prepare_inputs['sft_format'][0]}", answer)
-    det_pattern = r"<\|det\|>\[\[(.+)]]<\|\/det\|>"
-
-    det_match = re.search(det_pattern, answer)
-    if det_match is None:
-        return text_input, [], image
     
-    det_content = det_match.group(1)
-    bbox = [int(v.strip()) for v in det_content.split(",")]
-
-    scaled_boxes = rescale_bounding_boxes([bbox], image.width, image.height)
-    return answer, scaled_boxes, draw_bounding_boxes(image, scaled_boxes)
+    with torch.no_grad():
+
+        # run image encoder to get the image embeddings
+        inputs_embeds = vl_gpt.prepare_inputs_embeds(**prepare_inputs)
+
+        inputs_embeds, past_key_values = vl_gpt.incremental_prefilling(
+            input_ids=prepare_inputs.input_ids,
+            images=prepare_inputs.images,
+            images_seq_mask=prepare_inputs.images_seq_mask,
+            images_spatial_crop=prepare_inputs.images_spatial_crop,
+            attention_mask=prepare_inputs.attention_mask,
+            chunk_size=512 # prefilling size
+        )
+
+        # run the model to get the response
+        outputs = vl_gpt.generate(
+            inputs_embeds=inputs_embeds,
+            input_ids=prepare_inputs.input_ids,
+            images=prepare_inputs.images,
+            images_seq_mask=prepare_inputs.images_seq_mask,
+            images_spatial_crop=prepare_inputs.images_spatial_crop,
+            attention_mask=prepare_inputs.attention_mask,
+            past_key_values=past_key_values,
+            pad_token_id=tokenizer.eos_token_id,
+            bos_token_id=tokenizer.bos_token_id,
+            eos_token_id=tokenizer.eos_token_id,
+            max_new_tokens=512,
+            do_sample=False,
+            use_cache=True,
+        )
+
+        answer = tokenizer.decode(outputs[0][len(prepare_inputs.input_ids[0]):].cpu().tolist(), skip_special_tokens=False)
+        print(f"{prepare_inputs['sft_format'][0]}", answer)
+        det_pattern = r"<\|det\|>\[\[(.+)]]<\|\/det\|>"
+
+        det_match = re.search(det_pattern, answer)
+        if det_match is None:
+            return text_input, [], image
+        
+        det_content = det_match.group(1)
+        bbox = [int(v.strip()) for v in det_content.split(",")]
+
+        scaled_boxes = rescale_bounding_boxes([bbox], image.width, image.height)
+        return answer, scaled_boxes, draw_bounding_boxes(image, scaled_boxes)
 
 
 @spaces.GPU