[bugfix] fix keye_vl (#5848)

Author: Jintao-Huang

examples/models/keye/train.sh

@@ -0,0 +1,26 @@
+# 24GiB
+CUDA_VISIBLE_DEVICES=0 \
+swift sft \
+    --model Kwai-Keye/Keye-VL-1_5-8B \
+    --dataset 'AI-ModelScope/LaTeX_OCR:human_handwrite#20000' \
+    --split_dataset_ratio 0.01 \
+    --train_type lora \
+    --torch_dtype bfloat16 \
+    --num_train_epochs 1 \
+    --per_device_train_batch_size 1 \
+    --per_device_eval_batch_size 1 \
+    --learning_rate 1e-4 \
+    --lora_rank 8 \
+    --lora_alpha 32 \
+    --target_modules all-linear \
+    --freeze_vit true \
+    --gradient_accumulation_steps 16 \
+    --eval_steps 50 \
+    --save_steps 50 \
+    --save_total_limit 2 \
+    --logging_steps 5 \
+    --max_length 2048 \
+    --output_dir output \
+    --warmup_ratio 0.05 \
+    --dataset_num_proc 4 \
+    --dataloader_num_workers 4

swift/llm/model/model/qwen.py

@@ -685,6 +685,8 @@ def _new_read_video_decord(ele: dict):
         backends = getattr(vision_process, 'VIDEO_READER_BACKENDS', None)
         if isinstance(backends, dict):
             backends['decord'] = _new_read_video_decord
+        elif backends is None:  # keye_vl
+            vision_process._read_video_decord = _new_read_video_decord
     vision_process._patch = True
     return res
 

swift/llm/template/template/kwai.py

@@ -41,11 +41,12 @@ def replace_tag(self, media_type: Literal['image', 'video', 'audio'], index: int
             video = inputs.videos[index]
             if os.path.isdir(video):
                 video = [os.path.join(video, fname) for fname in os.listdir(video)]
-            video, video_kwargs = fetch_video({'video': video}, return_video_sample_fps=True)
+            video, video_kwargs = fetch_video({'video': video})
             if isinstance(video, torch.Tensor):
                 video = video.to(torch.uint8)
             inputs.videos[index] = video
-            inputs.mm_processor_kwargs.setdefault('fps', []).append(video_kwargs)
+            for k, v in video_kwargs.items():
+                inputs.mm_processor_kwargs.setdefault(k, []).append(v)
             return ['<|vision_start|><|video_pad|><|vision_end|>']
 
     def _encode(self, inputs: StdTemplateInputs) -> Dict[str, Any]:
@@ -62,25 +63,24 @@ def _encode(self, inputs: StdTemplateInputs) -> Dict[str, Any]:
                     media_inputs = processor.image_processor(images=mm_data, return_tensors='pt', do_resize=False)
                     media_grid_thw = media_inputs['image_grid_thw']
                 else:
-                    kwargs = {}
-                    if hasattr(processor, 'video_processor'):
-                        processor_func = processor.video_processor
-                    else:
-                        processor_func = processor.image_processor
-                        kwargs['images'] = None
-                    media_inputs = processor_func(videos=mm_data, return_tensors='pt', do_resize=False, **kwargs)
+                    split_token = self._tokenize('\n')[0]
+                    media_inputs = processor(
+                        text=['\n'.join(['<|video_pad|>'] * len(mm_data))],
+                        videos=mm_data,
+                        return_tensors='pt',
+                        **inputs.mm_processor_kwargs)
+                    splited_tokens = self._split_list(media_inputs['input_ids'][0].tolist(), split_token)
                     media_grid_thw = media_inputs['video_grid_thw']
                     media_token = self.video_token_id
-                    fps = inputs.mm_processor_kwargs['fps']
-                    media_inputs['second_per_grid_ts'] = [
-                        processor.image_processor.temporal_patch_size / tmp for tmp in fps
-                    ]
                 idx_list = findall(input_ids, media_token)
                 merge_length = processor.image_processor.merge_size**2
 
                 def _get_new_tokens(i):
-                    token_len = (media_grid_thw[i].prod() // merge_length)
-                    return [media_token] * token_len
+                    if media_type == 'images':
+                        token_len = (media_grid_thw[i].prod() // merge_length)
+                        return [media_token] * token_len
+                    else:
+                        return splited_tokens[i]
 
                 input_ids, labels, loss_scale = self._extend_tokens(input_ids, labels, loss_scale, idx_list,
                                                                     _get_new_tokens)
@@ -291,6 +291,14 @@ def _data_collator_mm_data(self, batch: List[Dict[str, Any]]) -> Dict[str, Any]:
 
 # Register the Keye VL template
 register_template(KeyeTemplateMeta(MLLMTemplateType.keye_vl, template_cls=KeyeVLTemplate))
+
+
+class KeyeVL1_5Template(KeyeVLTemplate):
+
+    def _post_encode(self, model, inputs: Dict[str, Any]) -> Dict[str, Any]:
+        return super(KeyeVLTemplate, self)._post_encode(model, inputs)
+
+
 register_template(
     KeyeTemplateMeta(
-        MLLMTemplateType.keye_vl_1_5, template_cls=KeyeVLTemplate, default_system='You are a helpful assistant.'))
+        MLLMTemplateType.keye_vl_1_5, template_cls=KeyeVL1_5Template, default_system='You are a helpful assistant.'))

tests/test_align/test_template/test_video.py

@@ -173,15 +173,25 @@ def test_glm4_5v():
 
 
 def test_keye_vl():
-    pt_engine = PtEngine('Kwai-Keye/Keye-VL-8B-Preview', attn_impl='flash_attention_2')
-    messages = [{'role': 'user', 'content': '<video>What happened in the video?'}]
+    pt_engine = PtEngine('Kwai-Keye/Keye-VL-8B-Preview')
+    messages = [{'role': 'user', 'content': '<video>Describe this video.'}]
     videos = ['https://modelscope-open.oss-cn-hangzhou.aliyuncs.com/images/baby.mp4']
     response = _infer_model(pt_engine, messages=messages, videos=videos)
     pt_engine.default_template.template_backend = 'jinja'
     response2 = _infer_model(pt_engine, messages=messages, videos=videos)
     assert response == response2
 
 
+def test_keye_vl_1_5():
+    pt_engine = PtEngine('Kwai-Keye/Keye-VL-1_5-8B')
+    messages = [{'role': 'user', 'content': '<video>Describe this video.'}]
+    videos = ['https://modelscope-open.oss-cn-hangzhou.aliyuncs.com/images/baby.mp4']
+    response = _infer_model(pt_engine, messages=messages, videos=videos)
+    assert response[:200] == ('The video features a young child sitting on a bed, engrossed in '
+                              'reading a book. The child is wearing a light blue sleeveless top and pink '
+                              'pants. The book appears to be a hardcover with illustrations, ')
+
+
 def test_ovis2_5():
     pt_engine = PtEngine('AIDC-AI/Ovis2.5-2B')
     messages = [{'role': 'user', 'content': '<video>Describe this video in detail.'}]
@@ -241,9 +251,10 @@ def test_minicpmv4_5():
     # test_qwen2_5_vl()
     # test_qwen2_5_omni()
     # test_glm4_1v()  # bug now, wait model fix
-    # test_keye_vl()
+    test_keye_vl()
+    test_keye_vl_1_5()
     # test_glm4_5v()
     # test_ovis2_5()
     # test_interns1()
     # test_internvl3_5()
-    test_minicpmv4_5()
+    # test_minicpmv4_5()

tests/test_align/test_template/test_vision.py

@@ -790,13 +790,13 @@ def test_minicpmv4_5():
     # test_kimi_vl_thinking()
     # test_glm4_1v()
     # test_gemma3n()
-    # test_keye_vl()
+    test_keye_vl()
     # test_dots_ocr()
     # test_glm4_5v()
     # test_interns1()
     # test_internvl3_5()
     # test_minicpmv4_5()
-    # test_keye_vl_1_5()
+    test_keye_vl_1_5()
     # test_internvl3_hf()
     # test_internvl3_5_hf()
-    test_internvl_gpt_hf()
+    # test_internvl_gpt_hf()