Add args support for fast image processors (#37018)

* add args support to fast image processors

* add comment for clarity

* fix-copies

* Handle child class args passed as both args or kwargs in call and preprocess functions

* revert support args passed as kwargs in overwritten preprocess

* fix image processor errors
This commit is contained in:
Yoni Gozlan
2025-05-16 12:01:46 -04:00
committed by GitHub
parent d69945e5fc
commit 0ba95564b7
12 changed files with 68 additions and 71 deletions

View File

@@ -19,8 +19,9 @@ import warnings
import numpy as np
import requests
from packaging import version
from transformers.testing_utils import is_flaky, require_torch, require_vision
from transformers.testing_utils import is_flaky, require_torch, require_torch_gpu, require_vision, slow, torch_device
from transformers.utils import is_torch_available, is_torchvision_available, is_vision_available
from ...test_image_processing_common import ImageProcessingTestMixin, prepare_image_inputs
@@ -334,3 +335,24 @@ class VitMatteImageProcessingTest(ImageProcessingTestMixin, unittest.TestCase):
self.assertLessEqual(
torch.mean(torch.abs(encoding_slow.pixel_values - encoding_fast.pixel_values)).item(), 1e-3
)
@slow
@require_torch_gpu
@require_vision
def test_can_compile_fast_image_processor(self):
# override as trimaps are needed for the image processor
if self.fast_image_processing_class is None:
self.skipTest("Skipping compilation test as fast image processor is not defined")
if version.parse(torch.__version__) < version.parse("2.3"):
self.skipTest(reason="This test requires torch >= 2.3 to run.")
torch.compiler.reset()
input_image = torch.randint(0, 255, (3, 224, 224), dtype=torch.uint8)
dummy_trimap = np.random.randint(0, 3, size=input_image.shape[1:])
image_processor = self.fast_image_processing_class(**self.image_processor_dict)
output_eager = image_processor(input_image, dummy_trimap, device=torch_device, return_tensors="pt")
image_processor = torch.compile(image_processor, mode="reduce-overhead")
output_compiled = image_processor(input_image, dummy_trimap, device=torch_device, return_tensors="pt")
torch.testing.assert_close(output_eager.pixel_values, output_compiled.pixel_values, rtol=1e-4, atol=1e-4)