[NicolasHug]

Context: I'm the maintainer for both TorchCodec and TorchVision

Motivation

JPEGs are decoded on the GPU by default (gpu_image_decode = True) with torchvision.io.decode_jpeg(device="cuda"). That decoder has three known bugs in sglang:

• It writes its output before queued GPU work has finished reading the memory ([Bug] GPU JPEG decode can overwrite memory that queued GPU work still reads #43075, CUDA decode_jpeg can overwrite memory that queued work on the current stream still reads pytorch/vision#9699). It allocates on the current stream but decodes on a private nvJPEG stream that doesn't wait for it. When the allocator hands it a block that queued kernels are still reading, those kernels see overwritten data. In the reported case another image's pixel_values became NaN and the request generated garbage.
• It returns grayscale JPEGs as [1, H, W], which then fails RGB normalization or torch.stack in the HF processors ([Bug] Grayscale JPEG's processing failes on normalization step #44748, [Bug] Gemma-4 mm: single non-RGB image crashes vision tower and kills scheduler (mat1 256 vs 768) #26751). Fixes proposed in [Fix] Normalize grayscale JPEGs in GPU image decoding #35101, [Fix] multimodal: decode JPEG to RGB on GPU path to fix grayscale crash #28828 and [Fix] multimodal: expand single-channel GPU JPEG tensors to RGB on load #44781 were never merged.
• It ignores EXIF orientation, so portrait phone photos reach the model rotated ([Bug] EXIF orientation is never applied when decoding images for VLM serving #33464). The PIL path doesn't apply EXIF orientation either (the fix, [Fix] multimodal: apply EXIF orientation when decoding images #33449, is still open), except for processors that opt into smart_rgb_conversion.

This PR switches the GPU path to torchcodec's decode_jpeg, which is already a dependency (pinned at 0.17.0). It decodes on the caller's current stream and synchronizes with work already queued there (meta-pytorch/torchcodec#1634), outputs RGB with mode="RGB", and applies EXIF orientation. The PIL path now applies EXIF orientation too, so the GPU and CPU paths agree.

This is a behavior change for images whose EXIF Orientation tag is 2–8 (in practice, most portrait phone photos): the model now sees them upright, as viewers display them. HF transformers' load_image and vLLM do the same. Images without such a tag are unaffected.

Fixes #43075
Fixes #33464
Fixes the GPU-decode case of #44748, the reported 1-channel tensor.
Addresses the grayscale-JPEG trigger of #26751.
Supersedes #43076, #33449 and #35101.

Modifications

• srt/utils/common.py:
  • _load_image decodes JPEGs on the GPU with torchcodec.decoders.decode_jpeg(image_bytes, mode="RGB", device="cuda").
  • If torchcodec is missing or fails, it warns once per distinct error and falls back to PIL.
  • PIL-decoded images and caller-supplied PIL images go through a new _apply_exif_orientation helper. It returns the image itself (no copy) when there's nothing to undo, doesn't rotate the caller's object in place, and treats malformed EXIF as no rotation.
  • load_image takes apply_exif_orientation (default True). With False, images keep their stored orientation on every path; since torchcodec always applies EXIF, such JPEGs are decoded with PIL instead.
  • smart_to_rgb no longer applies EXIF orientation itself, since images reach it already upright.
• Opt-out for Kimi-K3 and EmbeddingGemma2. Their GPU decoder (nvImageCodec, gpu_image_decode="nvjpeg_fancy") ignores EXIF orientation. They set apply_exif_orientation = False on their processors, and Kimi-K3 also sets it in its EncoderMediaProcessorConfig. As a result their PIL fallback and their non-JPEG inputs keep today's behavior too. The flag is passed through BaseMultimodalProcessor._load_single_item, MediaArtifactCacheMixin.decode_media_snapshot and the encoder-disaggregation preprocessor. It is also part of the preprocess-cache fingerprint, so artifacts cached before this change aren't reused.
• srt/multimodal/media_artifacts/base.py: PIL inputs on the preprocess-cache path go through load_image, so they get the same loading and orientation as everywhere else.
• python/pyproject.toml: torchcodec is now installed unconditionally. It used to be excluded on Linux ARM because no wheels were available there. torchcodec has shipped Linux aarch64 wheels since 0.11, and 0.17.0 has CUDA-enabled aarch64 wheels for every supported CPython. So GB200/GH200 keep GPU JPEG decoding. As a consequence, video and audio decoding on Linux ARM also use torchcodec, as they already do on x86 (VideoDecoderWrapper and load_audio use torchcodec whenever it imports). decord2 and PyAV stay installed there as fallbacks. The Docker images already install FFmpeg on every architecture.

Accuracy Tests

For JPEGs that are already RGB and have no EXIF rotation, the GPU path's output is unchanged. I compared torchcodec 0.17.0 against torchvision 0.29.1, both decoding on CUDA (RTX 4080, torch 2.14.1+cu130). On every RGB JPEG in torchcodec's test set and doc thumbnails (gradient, 15 real photos), the two outputs were bit-identical (torch.equal). Both differ from PIL by the same amount, since both use nvJPEG's default chroma upsampling. So I didn't run an end-to-end eval.

Model inputs change only in the cases this PR fixes: grayscale JPEGs (now 3-channel instead of an error) and photos with an EXIF orientation (now upright).

Speed Tests and Profiling

Median of 50 runs on a 1920x1280 JPEG (q=90, 4:2:0), RTX 4080:

	ms / image
PIL decode + CHW tensor + .cuda()	16.39
torchvision decode_jpeg (CUDA)	3.77
torchcodec decode_jpeg (CUDA)	3.78
torchcodec, batch of 8 (CUDA)	3.20

Per-image speed is unchanged. Batching a request's JPEGs into one nvJPEG call would save about 15%; that's left for a follow-up. The first-use lazy import costs about 930 ms for torchcodec against about 760 ms for torchvision, once per tokenizer process.

Tests

New tests in test/registered/unit/multimodal/test_base_processor_image_decode.py (CPU) and test/registered/unit/utils/test_common.py (1-GPU CUDA CI). Each non-regression test names its issue in its docstring. All of them fail on main and pass with this PR:

• Grayscale (CUDA, [Bug] Grayscale JPEG's processing failes on normalization step #44748): a grayscale JPEG decodes to (3, H, W).
• EXIF on the GPU path (CUDA, [Bug] EXIF orientation is never applied when decoding images for VLM serving #33464): a JPEG with Orientation=6 matches PIL's exif_transpose. The mean abs difference is about 1, against about 128 for the opposite rotation.
• Stream race (CUDA, [Bug] GPU JPEG decode can overwrite memory that queued GPU work still reads #43075): decoding into memory freed while kernels queued behind torch.cuda._sleep still read it. On main this corrupted those kernels' results in 7 of 20 runs on an RTX 4080; with torchcodec, 0 of 20.
• EXIF on the PIL path ([Bug] EXIF orientation is never applied when decoding images for VLM serving #33464): all 8 orientations across JPEG, PNG and WebP.
• Properties: caller-supplied PIL images are oriented without being rotated in place, unrotated images aren't copied, and malformed EXIF counts as no rotation.
• Fallbacks: a torchcodec decode error and a missing torchcodec both fall back to PIL.
• Opt-out: a processor with apply_exif_orientation = False gets the stored orientation, and torchcodec is never called for it.

Locally, those two files plus test_audio_container_decode.py, test_base_processor_bad_input.py and test_media_artifact_processor.py all pass, as does test_encode_server.py: 87 tests and 59 subtests. pre-commit is clean.

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

---

CI States

Latest PR Test (Base): ❌ Run #39990804200
Latest PR Test (Extra): ❌ Run #39990804122
Latest PR Test (AMD ROCm 10): ❌ Run #39990804063

[NicolasHug]

Hi @mickqian, pinging you here as your Slack profile asks. I'm Nicolas, on the PyTorch team at Meta. I maintain torchcodec and torchvision. I noticed a few areas for improvement in sglang's image/video decoding, and this is the first: it moves GPU JPEG decoding from torchvision to torchcodec, which fixes the stream race, grayscale and EXIF bugs (#43075, #44748, #33464).

Some other improvements I'd love to follow up on:

• CPU image decoding with torchcodec to tensors, which skips the PIL-to-tensor conversion and is generally ~10% faster than PIL
• Batching a request's JPEGs into one nvJPEG call: about 15% faster per image in my tests.

Happy to connect on Slack to chat more. Thanks!