fix bug for webui infer

Former-commit-id: 17768832908cc59ab64ed72522b2954c575ce21d
2024-10-16 01:09:33 +08:00
parent 67645c0db8
commit 4845a76535
2 changed files with 9 additions and 2 deletions
--- a/src/llamafactory/data/mm_plugin.py
+++ b/src/llamafactory/data/mm_plugin.py
@@ -513,6 +513,12 @@ class PixtralPlugin(BasePlugin):
    ) -> Dict[str, Union[List[int], "torch.Tensor"]]:
        self._validate_input(images, videos)
        mm_inputs = self._get_mm_inputs(images, videos, processor)
+        # hack for hf engine
+        if mm_inputs.get("pixel_values") and len(mm_inputs.get("pixel_values")[0]) == 1:
+            mm_inputs["pixel_values"] = mm_inputs["pixel_values"][0][0].unsqueeze(0)
+
+        if mm_inputs.get("image_sizes"):
+            del mm_inputs["image_sizes"]

        return mm_inputs