@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 165 lines 5.96 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 from io import BytesIO
5 from pathlib import Path
6
7 import numpy as np
8 import pybase64
9 import torch
10 from PIL import Image
11
12 import vllm.envs as envs
13 from vllm.utils.serial_utils import tensor2base64
14 from vllm.utils.sparse_utils import (
15 check_sparse_tensor_invariants_threadsafe,
16 safe_to_dense,
17 )
18
19 from ..image import convert_image_mode, normalize_image, rgba_to_rgb
20 from .base import MediaIO, MediaWithBytes
21
22 MAGIC_NUMPY_PREFIX = b"\x93NUMPY" # https://numpy.org/devdocs/reference/generated/numpy.lib.format.html#format-version-1-0
23
24
25 class ImageMediaIO(MediaIO[Image.Image]):
26 """Configuration values can be user-provided either by --media-io-kwargs or
27 by the runtime API field "media_io_kwargs". Ensure proper validation and
28 error handling.
29 """
30
31 def __init__(self, image_mode: str | None = "RGB", **kwargs) -> None:
32 super().__init__()
33
34 # Target mode for loaded images; `None` keeps the original mode
35 # (i.e. no conversion, alpha channel is preserved as-is).
36 self.image_mode = image_mode
37 # `kwargs` contains custom arguments from
38 # --media-io-kwargs for this modality, merged with
39 # per-request runtime media_io_kwargs via merge_kwargs().
40 # They can be passed to the underlying
41 # media loaders (e.g. custom implementations)
42 # for flexible control.
43 self.kwargs = kwargs
44
45 # Extract RGBA background color from kwargs if provided
46 # Default to white background for backward compatibility
47 rgba_bg = kwargs.get("rgba_background_color", (255, 255, 255))
48 # Convert list to tuple for consistency
49 if isinstance(rgba_bg, list):
50 rgba_bg = tuple(rgba_bg)
51
52 # Validate rgba_background_color format
53 if not (
54 isinstance(rgba_bg, tuple)
55 and len(rgba_bg) == 3
56 and all(isinstance(c, int) and 0 <= c <= 255 for c in rgba_bg)
57 ):
58 raise ValueError(
59 "rgba_background_color must be a list or tuple of 3 integers "
60 "in the range [0, 255]."
61 )
62 self.rgba_background_color = rgba_bg
63
64 def _convert_image_mode(
65 self, image: Image.Image | MediaWithBytes[Image.Image]
66 ) -> Image.Image:
67 """Convert image mode with custom background color."""
68 if isinstance(image, MediaWithBytes):
69 image = image.media
70 if self.image_mode is None or image.mode == self.image_mode:
71 return image
72 elif image.mode == "RGBA" and self.image_mode == "RGB":
73 return rgba_to_rgb(image, self.rgba_background_color)
74 else:
75 return convert_image_mode(
76 image, self.image_mode, self.rgba_background_color
77 )
78
79 def load_bytes(self, data: bytes) -> MediaWithBytes[Image.Image]:
80 try:
81 image = Image.open(BytesIO(data))
82 w, h = image.size
83 max_pixels = envs.VLLM_MAX_IMAGE_PIXELS
84 if max_pixels > 0 and w * h > max_pixels:
85 raise ValueError(
86 f"Image dimensions {w}x{h} ({w * h} pixels) exceed "
87 f"the maximum of {max_pixels} pixels. Set "
88 f"VLLM_MAX_IMAGE_PIXELS to increase this limit."
89 )
90 image = normalize_image(image)
91 image.load()
92 converted = self._convert_image_mode(image)
93 except (OSError, Image.UnidentifiedImageError) as e:
94 raise ValueError(f"Failed to load image: {e}") from e
95
96 io_config = None
97 if converted is not image:
98 io_config = {
99 "image_mode": self.image_mode,
100 "rgba_background_color": self.rgba_background_color,
101 }
102 return MediaWithBytes(converted, data, io_config)
103
104 def load_base64(self, media_type: str, data: str) -> MediaWithBytes[Image.Image]:
105 return self.load_bytes(pybase64.b64decode(data, validate=True))
106
107 def load_file(self, filepath: Path) -> MediaWithBytes[Image.Image]:
108 return self.load_bytes(filepath.read_bytes())
109
110 def encode_base64(
111 self,
112 media: Image.Image,
113 *,
114 image_format: str = "PNG",
115 ) -> str:
116 image = media
117
118 with BytesIO() as buffer:
119 image = self._convert_image_mode(image)
120 image.save(buffer, image_format)
121 data = buffer.getvalue()
122
123 return pybase64.b64encode(data).decode("utf-8")
124
125
126 class ImageEmbeddingMediaIO(MediaIO[torch.Tensor]):
127 """Image embedding MediaIO implementation.
128
129 Configuration values can be user-provided either by --media-io-kwargs or
130 by the runtime API field "media_io_kwargs". Ensure proper validation and
131 error handling.
132 """
133
134 def __init__(self) -> None:
135 super().__init__()
136
137 def _load_pickled_torch(self, data: bytes) -> torch.Tensor:
138 buffer = BytesIO(data)
139 with check_sparse_tensor_invariants_threadsafe():
140 tensor = torch.load(buffer, weights_only=True)
141 return safe_to_dense(tensor, parameter="image_embeds")
142
143 def _load_numpy(self, data: bytes) -> torch.Tensor:
144 with BytesIO(data) as buffer:
145 return torch.from_numpy(np.load(buffer))
146
147 def load_bytes(self, data: bytes) -> torch.Tensor:
148 if data[:6] == MAGIC_NUMPY_PREFIX:
149 return self._load_numpy(data)
150
151 return self._load_pickled_torch(data)
152
153 def load_base64(self, media_type: str, data: str) -> torch.Tensor:
154 return self.load_bytes(pybase64.b64decode(data, validate=True))
155
156 def load_file(self, filepath: Path) -> torch.Tensor:
157 if filepath.suffix == ".npy":
158 return torch.from_numpy(np.load(filepath))
159
160 with check_sparse_tensor_invariants_threadsafe():
161 tensor = torch.load(filepath, weights_only=True)
162 return safe_to_dense(tensor, parameter="image_embeds")
163
164 def encode_base64(self, media: torch.Tensor) -> str:
165 return tensor2base64(media)