From 9db963aeed419c8379c6d6425186fec0bfb86908 Mon Sep 17 00:00:00 2001 From: Dominik Niedziela <99881522+dom-dziela@users.noreply.github.com> Date: Wed, 18 Sep 2024 17:38:31 +0200 Subject: [PATCH] enforce original size to be a list (#33564) * enforce original size to be a list * formatting * apply datatype change to unpad_image in llava_next --- src/transformers/models/llava_next/modeling_llava_next.py | 6 ++++++ .../models/llava_onevision/modeling_llava_onevision.py | 6 ++++++ 2 files changed, 12 insertions(+) diff --git a/src/transformers/models/llava_next/modeling_llava_next.py b/src/transformers/models/llava_next/modeling_llava_next.py index c1d1ca8c276d7a..bf76921090b244 100644 --- a/src/transformers/models/llava_next/modeling_llava_next.py +++ b/src/transformers/models/llava_next/modeling_llava_next.py @@ -123,6 +123,12 @@ def unpad_image(tensor, original_size): Returns: `torch.Tensor`: The unpadded image tensor. """ + if not isinstance(original_size, (list, tuple)): + if not isinstance(original_size, (torch.Tensor, np.ndarray)): + raise TypeError( + f"image_size invalid type: {type(original_size)} not valid, should be either list, tuple, np.ndarray or tensor" + ) + original_size = original_size.tolist() original_height, original_width = original_size current_height, current_width = tensor.shape[1:] diff --git a/src/transformers/models/llava_onevision/modeling_llava_onevision.py b/src/transformers/models/llava_onevision/modeling_llava_onevision.py index 697ea84fea5040..d3200fb5193d4b 100644 --- a/src/transformers/models/llava_onevision/modeling_llava_onevision.py +++ b/src/transformers/models/llava_onevision/modeling_llava_onevision.py @@ -124,6 +124,12 @@ def unpad_image(tensor, original_size): Returns: `torch.Tensor`: The unpadded image tensor. """ + if not isinstance(original_size, (list, tuple)): + if not isinstance(original_size, (torch.Tensor, np.ndarray)): + raise TypeError( + f"image_size invalid type: {type(original_size)} not valid, should be either list, tuple, np.ndarray or tensor" + ) + original_size = original_size.tolist() original_height, original_width = original_size current_height, current_width = tensor.shape[1:]