支持多图和任意分辨率

qwen2_tokenizer_images.py

@@ -79,14 +79,15 @@ class Tokenizer_Http():
 
     def encode(self, content):
         text = [f'<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n<|im_start|>user\n{content}<|im_end|>\n<|im_start|>assistant\n']
-        
         input_ids = self.tokenizer(text)
         return input_ids["input_ids"][0]
 
-    def encode_vpm(self, content="Describe this image."):
+    def encode_vpm(self, content="Describe this image.", num_img=1, img_token_num=256):
 
         # official implementation 
-        text = f'<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n<|im_start|>user\n<|vision_start|>' + '<|image_pad|>' * 256 + f'<|vision_end|>{content}<|im_end|>\n<|im_start|>assistant\n'
+        imgs_token = '<|vision_start|>' +  '<|image_pad|>'*img_token_num + '<|vision_end|>'
+        imgs_token *= num_img
+        text = f'<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n<|im_start|>user\n{imgs_token}{content}<|im_end|>\n<|im_start|>assistant\n'
         
         output_kwargs = {'text_kwargs': {'padding': True, 'return_tensors': 'pt'}, 'images_kwargs': {'return_tensors': 'pt'}, 'audio_kwargs': {'padding': True, 'return_tensors': 'pt'}, 'videos_kwargs': {'fps': 2.0, 'return_tensors': 'pt'}, 'common_kwargs': {'return_tensors': 'pt'}}
         
@@ -205,7 +206,7 @@ class Request(BaseHTTPRequestHandler):
             if 'img_prompt' in req:
                 b_img_prompt = req['img_prompt']
             if b_img_prompt:
-                token_ids = tokenizer.encode_vpm(prompt)
+                token_ids = tokenizer.encode_vpm(prompt, req["num_img"], req["img_token_num"])
             else:
                 token_ids = tokenizer.encode(prompt)