Ensure that `OneFormerProcessor` place text `task_inputs` to the same device as other inputs

[simonreise]

System Info

• transformers version: 5.0.0.dev0
• Platform: Windows-10-10.0.26200-SP0
• Python version: 3.11.9
• Huggingface_hub version: 1.2.1
• Safetensors version: 0.5.3
• Accelerate version: not installed
• Accelerate config: not found
• DeepSpeed version: not installed
• PyTorch version (accelerator?): 2.9.1+cu128 (CUDA)
• Using distributed or parallel set-up in script?:
• Using GPU in script?:
• GPU type: NVIDIA GeForce RTX 2060

Who can help?

@yonigozlan @molbap

Information

• The official example scripts
• My own modified scripts

Tasks

• An officially supported task in the examples folder (such as GLUE/SQuAD, ...)
• My own task or dataset (give details below)

Reproduction

Fast image processors can keep the output tensors on the same device as the input tensors.

This works for most of the processors, but OneFormerProcessor, which can combine OneFormerImageProcessorFast and a tokenizer and takes both image and text as inputs, do not ensure that both outputs are on the same device.

Let's create a simple script that loads an image, moves it to cuda and then tries to preprocess it with OneFormerProcessor

import transformers
from PIL import Image
import requests
from torchvision import transforms

to_tensor_transform = transforms.ToTensor()

processor = transformers.OneFormerImageProcessorFast()
processor = transformers.OneFormerProcessor(
    image_processor=processor,
    tokenizer=transformers.AutoTokenizer.from_pretrained("openai/clip-vit-base-patch32"),
)

url = "https://huggingface.co/datasets/hf-internal-testing/fixtures_ade20k/resolve/main/ADE_val_00000001.jpg"
image = Image.open(requests.get(url, stream=True).raw)
# Convert image to tensor and move it to cuda
image = to_tensor_transform(image).to("cuda")

# Semantic Segmentation
inputs = processor(image, ["semantic"], return_tensors="pt")

When we check the inputs generated by the processor, we will see that pixel_values are on cuda, but task_inputs are not

Expected behavior

Probably OneFormerProcessor should not only tokenize the text inputs, but also move them to the same device as the images

[Rocketknight1]

cc @molbap @yonigozlan @zucchini-nlp - this does sound like a real bug to me!

[i3hz]

Hi I modified the __call__ method in src/transformers/models/oneformer/processing_oneformer.py and added a small snippet to check the device of the pixel_values parameter and then manually moving the task_inputs to the device .

if "pixel_values" in encoded_inputs:
            pixel_values = encoded_inputs["pixel_values"]
            device = None
            if isinstance(pixel_values, torch.Tensor):
                device = pixel_values.device
            elif isinstance(pixel_values, (list, tuple)) and len(pixel_values) > 0 and isinstance(pixel_values[0], torch.Tensor):
                device = pixel_values[0].device
            
            if device is not None:
                if "task_inputs" in encoded_inputs and isinstance(encoded_inputs["task_inputs"], torch.Tensor):
                    encoded_inputs["task_inputs"] = encoded_inputs["task_inputs"].to(device)
                if "text_inputs" in encoded_inputs and isinstance(encoded_inputs["text_inputs"], torch.Tensor):
                    encoded_inputs["text_inputs"] = encoded_inputs["text_inputs"].to(device)

Is this the expected fix ? If it is , I can submit a PR for it

Though it does add a little bit of latency(~0.1 ms) , which is probably expected for moving data from cpu to gpu .

[zucchini-nlp]

Looks more like a feature request to me. We currently do not accept a device in tokenizers, audio processors and slow image processors, so when an user passes device="cuda" to the processor call, the inputs will be in different device depending on modality

In general, I agree that this is not intuitive and returning all tensors on the same device is the best option. Though we might want to apply it for text models as well, such as tokenizer("hello world", return_tensors="pt", device="cuda") works. I need your opinion on it @Rocketknight1 @molbap @yonigozlan

As for the above example, imo it's not worth trying to move all inputs to the same device as the passed pixels. That means we need to change each processor's call method, doesn't fit well with many modalities and feels like hidden internal magic which I prefer to avoid