[WIP] Processor moves to device in __call__

src/transformers/feature_extraction_utils.py

@@ -44,6 +44,8 @@
 
 
 if TYPE_CHECKING:
+    import torch.device
+
     from .feature_extraction_sequence_utils import SequenceFeatureExtractor
 
 
@@ -70,16 +72,21 @@ class BatchFeature(UserDict):
             initialization.
         skip_tensor_conversion (`list[str]` or `set[str]`, *optional*):
             List or set of keys that should NOT be converted to tensors, even when `tensor_type` is specified.
+        device (`Union[torch.device, str]`, *optional*):
+            The device to create tensots on when tensor_type is "pt".
     """
 
     def __init__(
         self,
         data: Optional[dict[str, Any]] = None,
         tensor_type: Union[None, str, TensorType] = None,
         skip_tensor_conversion: Optional[Union[list[str], set[str]]] = None,
+        device: Optional[Union["torch.device", str]] = None,
     ):
         super().__init__(data)
-        self.convert_to_tensors(tensor_type=tensor_type, skip_tensor_conversion=skip_tensor_conversion)
+        if tensor_type != "pt" and device is not None:
+            raise ValueError(f"Inputs cannot be moved to a device={device} when `tensor_type` is not Pytorch tensors")
+        self.convert_to_tensors(tensor_type=tensor_type, skip_tensor_conversion=skip_tensor_conversion, device=device)
 
     def __getitem__(self, item: str) -> Any:
         """
@@ -117,13 +124,13 @@ def _get_is_as_tensor_fns(self, tensor_type: Optional[Union[str, TensorType]] =
                 raise ImportError("Unable to convert output to PyTorch tensors format, PyTorch is not installed.")
             import torch
 
-            def as_tensor(value):
+            def as_tensor(value, device=None):
                 if torch.is_tensor(value):
-                    return value
+                    return value.to(device=device)
 
                 # stack list of tensors if tensor_type is PyTorch (# torch.tensor() does not support list of tensors)
                 if isinstance(value, (list, tuple)) and len(value) > 0 and torch.is_tensor(value[0]):
-                    return torch.stack(value)
+                    return torch.stack(value).to(device=device)
 
                 # convert list of numpy arrays to numpy array (stack) if tensor_type is Numpy
                 if isinstance(value, (list, tuple)) and len(value) > 0:
@@ -136,14 +143,14 @@ def as_tensor(value):
                     ):
                         value = np.array(value)
                 if isinstance(value, np.ndarray):
-                    return torch.from_numpy(value)
+                    return torch.from_numpy(value).to(device=device)
                 else:
-                    return torch.tensor(value)
+                    return torch.tensor(value, device=device)
 
             is_tensor = torch.is_tensor
         else:
 
-            def as_tensor(value, dtype=None):
+            def as_tensor(value, dtype=None, device=None):
                 if isinstance(value, (list, tuple)) and isinstance(value[0], (list, tuple, np.ndarray)):
                     value_lens = [len(val) for val in value]
                     if len(set(value_lens)) > 1 and dtype is None:
@@ -158,6 +165,7 @@ def convert_to_tensors(
         self,
         tensor_type: Optional[Union[str, TensorType]] = None,
         skip_tensor_conversion: Optional[Union[list[str], set[str]]] = None,
+        device: Optional[Union["torch.device", str]] = None,
     ):
         """
         Convert the inner content to tensors.
@@ -168,6 +176,8 @@ def convert_to_tensors(
                 `None`, no modification is done.
             skip_tensor_conversion (`list[str]` or `set[str]`, *optional*):
                 List or set of keys that should NOT be converted to tensors, even when `tensor_type` is specified.
+            device (`Union[torch.device, str]`, *optional*):
+                The device to create tensots on when tensor_type is "pt".
 
         Note:
             Values that don't have an array-like structure (e.g., strings, dicts, lists of strings) are
@@ -191,7 +201,7 @@ def convert_to_tensors(
 
             try:
                 if not is_tensor(value):
-                    tensor = as_tensor(value)
+                    tensor = as_tensor(value, device=device)
                     self[key] = tensor
             except Exception as e:
                 if key == "overflowing_values":

src/transformers/image_processing_utils_fast.py

@@ -852,7 +852,6 @@ def preprocess(self, images: ImageInput, *args, **kwargs: Unpack[ImagesKwargs])
         # Extract parameters that are only used for preparing the input images
         do_convert_rgb = kwargs.pop("do_convert_rgb")
         input_data_format = kwargs.pop("input_data_format")
-        device = kwargs.pop("device")
 
         # Update kwargs that need further processing before being validated
         kwargs = self._further_process_kwargs(**kwargs)
@@ -864,7 +863,7 @@ def preprocess(self, images: ImageInput, *args, **kwargs: Unpack[ImagesKwargs])
         kwargs.pop("data_format")
 
         return self._preprocess_image_like_inputs(
-            images, *args, do_convert_rgb=do_convert_rgb, input_data_format=input_data_format, device=device, **kwargs
+            images, *args, do_convert_rgb=do_convert_rgb, input_data_format=input_data_format, **kwargs
         )
 
     def _preprocess_image_like_inputs(
@@ -885,7 +884,7 @@ def _preprocess_image_like_inputs(
         images = self._prepare_image_like_inputs(
             images=images, do_convert_rgb=do_convert_rgb, input_data_format=input_data_format, device=device
         )
-        return self._preprocess(images, *args, **kwargs)
+        return self._preprocess(images, *args, device=device, **kwargs)
 
     def _preprocess(
         self,
@@ -904,6 +903,7 @@ def _preprocess(
         pad_size: Optional[SizeDict],
         disable_grouping: Optional[bool],
         return_tensors: Optional[Union[str, TensorType]],
+        device: torch.device | str | None,
         **kwargs,
     ) -> BatchFeature:
         # Group images by size for batched resizing
@@ -932,7 +932,7 @@ def _preprocess(
         if do_pad:
             processed_images = self.pad(processed_images, pad_size=pad_size, disable_grouping=disable_grouping)
 
-        return BatchFeature(data={"pixel_values": processed_images}, tensor_type=return_tensors)
+        return BatchFeature(data={"pixel_values": processed_images}, tensor_type=return_tensors, device=device)
 
     def to_dict(self):
         encoder_dict = super().to_dict()

src/transformers/models/llava/image_processing_llava.py

@@ -14,7 +14,7 @@
 # limitations under the License.
 """Image processor class for LLaVa."""
 
-from typing import Optional, Union
+from typing import TYPE_CHECKING, Optional, Union
 
 import numpy as np
 
@@ -45,6 +45,8 @@
 
 logger = logging.get_logger(__name__)
 
+if TYPE_CHECKING:
+    import torch.device
 
 if is_vision_available():
     import PIL
@@ -294,6 +296,7 @@ def preprocess(
         return_tensors: Optional[Union[str, TensorType]] = None,
         data_format: Optional[ChannelDimension] = ChannelDimension.FIRST,
         input_data_format: Optional[Union[str, ChannelDimension]] = None,
+        device: Optional[Union["torch.device", str]] = None,
         **kwargs,
     ) -> PIL.Image.Image:
         """
@@ -426,7 +429,7 @@ def preprocess(
             image = to_channel_dimension_format(image, data_format, input_channel_dim=input_data_format)
             processed_images.append(image)
 
-        return BatchFeature(data={"pixel_values": processed_images}, tensor_type=return_tensors)
+        return BatchFeature(data={"pixel_values": processed_images}, tensor_type=return_tensors, device=device)
 
 
 __all__ = ["LlavaImageProcessor"]

src/transformers/models/llava/image_processing_llava_fast.py

@@ -112,6 +112,7 @@ def _preprocess(
         image_std: Optional[Union[float, list[float]]],
         disable_grouping: Optional[bool],
         return_tensors: Optional[Union[str, TensorType]],
+        device: torch.device | str | None,
         **kwargs,
     ) -> BatchFeature:
         # Group images by size for batched resizing
@@ -150,7 +151,7 @@ def _preprocess(
 
         processed_images = reorder_images(processed_images_grouped, grouped_images_index)
 
-        return BatchFeature(data={"pixel_values": processed_images}, tensor_type=return_tensors)
+        return BatchFeature(data={"pixel_values": processed_images}, tensor_type=return_tensors, device=device)
 
 
 __all__ = ["LlavaImageProcessorFast"]

src/transformers/models/llava/processing_llava.py

@@ -156,6 +156,7 @@ def __call__(
                 prompt_strings.append(sample)
 
         return_tensors = output_kwargs["text_kwargs"].pop("return_tensors", None)
+        device = output_kwargs["text_kwargs"].get("device")
         return_mm_token_type_ids = output_kwargs["text_kwargs"].pop("return_mm_token_type_ids", False)
         text_inputs = self.tokenizer(prompt_strings, **output_kwargs["text_kwargs"], return_tensors=None)
         self._check_special_mm_tokens(prompt_strings, text_inputs, modalities=["image"])
@@ -166,7 +167,11 @@ def __call__(
             mm_token_type_ids[array_ids == self.image_token_id] = 1
             text_inputs["mm_token_type_ids"] = mm_token_type_ids.tolist()
 
-        return BatchFeature(data={**text_inputs, **image_inputs}, tensor_type=return_tensors)
+        return BatchFeature(
+            data={**text_inputs, **image_inputs},
+            tensor_type=return_tensors,
+            device=device,
+        )
 
     def _get_num_multimodal_tokens(self, image_sizes=None, **kwargs):
         """

src/transformers/processing_utils.py

@@ -197,6 +197,8 @@ class TextKwargs(TypedDict, total=False):
             If set, will return tensors of a particular framework. Acceptable values are:
             - `'pt'`: Return PyTorch `torch.Tensor` objects.
             - `'np'`: Return NumPy `np.ndarray` objects.
+        device (`Union[str, torch.Tensor]`, *optional*):
+            The device to use for processing (e.g. "cpu", "cuda"), only relevant for fast image processing.
     """
 
     text_pair: Optional[Union[TextInput, PreTokenizedInput, list[TextInput], list[PreTokenizedInput]]]
@@ -219,6 +221,7 @@ class TextKwargs(TypedDict, total=False):
     padding_side: Optional[Literal["left", "right"]]
     return_mm_token_type_ids: Optional[bool]
     return_tensors: Annotated[Optional[Union[str, TensorType]], tensor_type_validator()]
+    device: Annotated[Optional[Union[str, "torch.device"]], device_validator()]
 
 
 class ImagesKwargs(TypedDict, total=False):
@@ -402,6 +405,8 @@ class AudioKwargs(TypedDict, total=False):
             If set, will return tensors of a particular framework. Acceptable values are:
             - `'pt'`: Return PyTorch `torch.Tensor` objects.
             - `'np'`: Return NumPy `np.ndarray` objects.
+        device (`Union[str, torch.Tensor]`, *optional*):
+            The device to use for processing (e.g. "cpu", "cuda"), only relevant for fast image processing.
     """
 
     sampling_rate: Annotated[Optional[int], positive_int()]
@@ -412,6 +417,7 @@ class AudioKwargs(TypedDict, total=False):
     pad_to_multiple_of: Annotated[Optional[int], positive_int()]
     return_attention_mask: Optional[bool]
     return_tensors: Annotated[Optional[Union[str, TensorType]], tensor_type_validator()]
+    device: Annotated[Optional[Union[str, "torch.device"]], device_validator()]
 
 
 class ProcessingKwargs(TypedDict, total=False):

src/transformers/tokenization_mistral_common.py

@@ -18,7 +18,7 @@
 from collections.abc import Callable, Mapping, Sized
 from enum import Enum
 from pathlib import Path
-from typing import Any, Union, overload
+from typing import TYPE_CHECKING, Any, Optional, Union, overload
 
 import numpy as np
 from huggingface_hub import create_repo
@@ -39,6 +39,10 @@
 from transformers.utils.import_utils import is_mistral_common_available, is_torch_available, requires
 
 
+if TYPE_CHECKING:
+    import torch.device
+
+
 if is_mistral_common_available():
     from mistral_common.protocol.instruct.request import ChatCompletionRequest
     from mistral_common.protocol.instruct.validator import ValidationMode
@@ -745,6 +749,7 @@ def _encode_plus(
         return_special_tokens_mask: bool = False,
         return_length: bool = False,
         verbose: bool = True,
+        device: Optional[Union["torch.device", str]] = None,
     ) -> BatchEncoding:
         def get_input_ids(text):
             if isinstance(text, str):
@@ -943,6 +948,7 @@ def prepare_for_model(
         return_length: bool = False,
         verbose: bool = True,
         prepend_batch_axis: bool = False,
+        device: Optional[Union["torch.device", str]] = None,
         **kwargs,
     ) -> BatchEncoding:
         """
@@ -1018,7 +1024,7 @@ def prepare_for_model(
             encoded_inputs["length"] = len(encoded_inputs["input_ids"])
 
         batch_outputs = BatchEncoding(
-            encoded_inputs, tensor_type=return_tensors, prepend_batch_axis=prepend_batch_axis
+            encoded_inputs, tensor_type=return_tensors, prepend_batch_axis=prepend_batch_axis, device=device
         )
 
         return batch_outputs

src/transformers/tokenization_python.py

@@ -19,7 +19,7 @@
 import bisect
 import unicodedata
 from collections import OrderedDict
-from typing import Any, overload
+from typing import TYPE_CHECKING, Any, Optional, Union, overload
 
 from .tokenization_utils_base import (
     INIT_TOKENIZER_DOCSTRING,
@@ -34,6 +34,9 @@
 from .utils import PaddingStrategy, TensorType, add_end_docstrings, logging
 
 
+if TYPE_CHECKING:
+    import torch.device
+
 logger = logging.get_logger(__name__)
 
 # Slow tokenizers are saved in a vocabulary plus three separated files
@@ -711,6 +714,7 @@ def _encode_plus(
         return_special_tokens_mask: bool = False,
         return_length: bool = False,
         verbose: bool = True,
+        device: Optional[Union["torch.device", str]] = None,
         **kwargs,
     ) -> BatchEncoding:
         # Detect batched inputs (list of sequences)
@@ -789,7 +793,7 @@ def _encode_plus(
                 return_attention_mask=return_attention_mask,
             )
 
-            return BatchEncoding(batch_outputs, tensor_type=return_tensors)
+            return BatchEncoding(batch_outputs, tensor_type=return_tensors, device=device)
 
         # Single sequence handling
         def get_input_ids(text):