Migrate Mistral3ImagePixelInputs to TensorSchema #21945
Conversation
|
👋 Hi! Thank you for contributing to the vLLM project. 💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels. Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging. To run CI, PR reviewers can either: Add 🚀 |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request successfully migrates Mistral3ImagePixelInputs to use TensorSchema for improved input validation, which is a great step towards better maintainability and robustness. I've identified a potential correctness issue where the new validation might fail for lists of images with varying sizes, a valid use case for this model. I've provided suggestions to address this by making the height and width dimensions dynamic and conditionally applying the bindings. Addressing this will ensure the validation logic is robust for all supported input formats.
There was a problem hiding this comment.
The TensorShape for pixel_values should account for the fact that when a list of images is provided, their heights and widths can vary. The docstring and usage in _process_image_input confirm this. Without marking h and w as dynamic dimensions, the TensorSchema validation will incorrectly fail for lists of images with different sizes.
By adding dynamic_dims={"h", "w"}, you allow TensorSchema to correctly handle lists of tensors with varying dimensions for height and width, which aligns with the model's capabilities.
| pixel_values: Annotated[ | |
| Union[torch.Tensor, list[torch.Tensor]], | |
| TensorShape("bn", 3, "h", "w"), | |
| ] | |
| pixel_values: Annotated[ | |
| Union[torch.Tensor, list[torch.Tensor]], | |
| TensorShape("bn", 3, "h", "w", dynamic_dims={"h", "w"}), | |
| ] |
There was a problem hiding this comment.
This suggestion seems valid, so I've updated to set "h" and "w" as dynamic dims.
Signed-off-by: Benji Beck <benjibeck@meta.com>
Signed-off-by: Benji Beck <benjibeck@meta.com>
DarkLight1337
added
the
ready
Signed-off-by: Benji Beck <benjibeck@meta.com> Co-authored-by: Cyrus Leung <tlleungac@connect.ust.hk> Signed-off-by: Duncan Moss <djm.moss@gmail.com>
Signed-off-by: Benji Beck <benjibeck@meta.com> Co-authored-by: Cyrus Leung <tlleungac@connect.ust.hk>
Signed-off-by: Benji Beck <benjibeck@meta.com> Co-authored-by: Cyrus Leung <tlleungac@connect.ust.hk>
Signed-off-by: Benji Beck <benjibeck@meta.com> Co-authored-by: Cyrus Leung <tlleungac@connect.ust.hk> Signed-off-by: Xiao Yu <xiao.yu@amd.com>
Signed-off-by: Benji Beck <benjibeck@meta.com> Co-authored-by: Cyrus Leung <tlleungac@connect.ust.hk>
Signed-off-by: Benji Beck <benjibeck@meta.com> Co-authored-by: Cyrus Leung <tlleungac@connect.ust.hk>
Signed-off-by: Benji Beck <benjibeck@meta.com> Co-authored-by: Cyrus Leung <tlleungac@connect.ust.hk>
Signed-off-by: Benji Beck <benjibeck@meta.com> Co-authored-by: Cyrus Leung <tlleungac@connect.ust.hk>
2 participants
Purpose
This PR migrates Mistral3ImagePixelInputs from a TypedDict-based definition to a structured TensorSchema model with runtime shape validation. This brings it in line with recent changes to Phi3VImagePixelInputs, and is part of a broader effort to improve input contract enforcement and debug-ability across multi-modal models.
Test Plan
Confirm validation works via standalone tests in tests/standalone_test/test_tensor_schema.py and rely on CI to check integration.
Test Result