diff --git a/lightx2v/models/input_encoders/hf/qwen25/qwen25_vlforconditionalgeneration.py b/lightx2v/models/input_encoders/hf/qwen25/qwen25_vlforconditionalgeneration.py index eb95f7fab..364753133 100755 --- a/lightx2v/models/input_encoders/hf/qwen25/qwen25_vlforconditionalgeneration.py +++ b/lightx2v/models/input_encoders/hf/qwen25/qwen25_vlforconditionalgeneration.py @@ -206,13 +206,16 @@ def infer(self, text, image_list=None): return_tensors="pt", ).to(AI_DEVICE) - encoder_hidden_states = self.text_encoder( - input_ids=model_inputs.input_ids, - attention_mask=model_inputs.attention_mask, - pixel_values=model_inputs.pixel_values, - image_grid_thw=model_inputs.image_grid_thw, - output_hidden_states=True, - ) + text_encoder_inputs = { + "input_ids": model_inputs.input_ids, + "attention_mask": model_inputs.attention_mask, + "pixel_values": model_inputs.pixel_values, + "image_grid_thw": model_inputs.image_grid_thw, + "output_hidden_states": True, + } + if "mm_token_type_ids" in model_inputs: + text_encoder_inputs["mm_token_type_ids"] = model_inputs.mm_token_type_ids + encoder_hidden_states = self.text_encoder(**text_encoder_inputs) else: template = self.prompt_template_encode drop_idx = self.prompt_template_encode_start_idx diff --git a/tests/test_qwen25_vl_input_forwarding.py b/tests/test_qwen25_vl_input_forwarding.py new file mode 100644 index 000000000..3cfdfb46c --- /dev/null +++ b/tests/test_qwen25_vl_input_forwarding.py @@ -0,0 +1,76 @@ +import unittest +from types import SimpleNamespace + +import torch + +from lightx2v.models.input_encoders.hf.qwen25 import qwen25_vlforconditionalgeneration as qwen25_vl + + +class _Batch(dict[str, torch.Tensor]): + def __getattr__(self, name): + return self[name] + + def to(self, device): + return self + + +class _Processor: + def __init__(self, model_inputs): + self.model_inputs = model_inputs + + def __call__(self, **kwargs): + return self.model_inputs + + +class _TextEncoder: + def __init__(self): + self.calls = [] + + def __call__(self, **kwargs): + self.calls.append(kwargs) + return SimpleNamespace(hidden_states=[torch.zeros((1, 3, 2))]) + + +class Qwen25VLInputForwardingTest(unittest.TestCase): + def test_infer_forwards_processor_multimodal_token_types(self): + mm_token_type_ids = torch.tensor([[0, 1, 0]]) + model_inputs = _Batch( + input_ids=torch.tensor([[1, 2, 3]]), + attention_mask=torch.ones((1, 3), dtype=torch.long), + pixel_values=torch.zeros((1, 3)), + image_grid_thw=torch.tensor([[1, 1, 1]]), + mm_token_type_ids=mm_token_type_ids, + ) + encoder = object.__new__(qwen25_vl.Qwen25_VLForConditionalGeneration_TextEncoder) + encoder.cpu_offload = False + encoder.is_layered = False + encoder.USE_IMAGE_ID_IN_PROMPT = True + encoder.config = {"task": "i2i"} + encoder.prompt_template_encode = "{}" + encoder.prompt_template_encode_start_idx = 0 + encoder.dtype = torch.float32 + encoder.processor = _Processor(model_inputs) + encoder.text_encoder = _TextEncoder() + encoder.preprocess_image = lambda image: (image, image, (1, 1), (1, 1)) + + previous_device = qwen25_vl.AI_DEVICE + qwen25_vl.AI_DEVICE = "cpu" + try: + encoder.infer(["prompt"], [object()]) + finally: + qwen25_vl.AI_DEVICE = previous_device + + self.assertIs(encoder.text_encoder.calls[0]["mm_token_type_ids"], mm_token_type_ids) + + model_inputs.pop("mm_token_type_ids") + qwen25_vl.AI_DEVICE = "cpu" + try: + encoder.infer(["prompt"], [object()]) + finally: + qwen25_vl.AI_DEVICE = previous_device + + self.assertNotIn("mm_token_type_ids", encoder.text_encoder.calls[1]) + + +if __name__ == "__main__": + unittest.main()