50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217 | class TransformerModel(RoboflowInferenceModel):
task_type = "lmm"
transformers_class = AutoModel
processor_class = AutoProcessor
default_dtype = torch.float16
generation_includes_input = False
needs_hf_token = False
skip_special_tokens = True
def __init__(
self, model_id, *args, dtype=None, huggingface_token=HUGGINGFACE_TOKEN, **kwargs
):
super().__init__(model_id, *args, **kwargs)
self.huggingface_token = huggingface_token
if self.needs_hf_token and self.huggingface_token is None:
raise RuntimeError(
"Must set environment variable HUGGINGFACE_TOKEN to load LoRA "
"(or pass huggingface_token to this __init__)"
)
self.dtype = dtype
if self.dtype is None:
self.dtype = self.default_dtype
self.cache_model_artefacts()
self.cache_dir = os.path.join(MODEL_CACHE_DIR, self.endpoint + "/")
self.initialize_model()
def initialize_model(self):
self.model = (
self.transformers_class.from_pretrained(
self.cache_dir,
device_map=DEVICE,
token=self.huggingface_token,
)
.eval()
.to(self.dtype)
)
self.processor = self.processor_class.from_pretrained(
self.cache_dir, token=self.huggingface_token
)
def preprocess(
self, image: Any, **kwargs
) -> Tuple[Image.Image, PreprocessReturnMetadata]:
pil_image = Image.fromarray(load_image_rgb(image))
image_dims = pil_image.size
return pil_image, PreprocessReturnMetadata({"image_dims": image_dims})
def postprocess(
self,
predictions: Tuple[str],
preprocess_return_metadata: PreprocessReturnMetadata,
**kwargs,
) -> LMMInferenceResponse:
text = predictions[0]
image_dims = preprocess_return_metadata["image_dims"]
response = LMMInferenceResponse(
response=text,
image=InferenceResponseImage(width=image_dims[0], height=image_dims[1]),
)
return [response]
def predict(self, image_in: Image.Image, prompt="", history=None, **kwargs):
model_inputs = self.processor(
text=prompt, images=image_in, return_tensors="pt"
).to(self.model.device)
input_len = model_inputs["input_ids"].shape[-1]
with torch.inference_mode():
prepared_inputs = self.prepare_generation_params(
preprocessed_inputs=model_inputs
)
generation = self.model.generate(
**prepared_inputs,
max_new_tokens=1000,
do_sample=False,
early_stopping=False,
no_repeat_ngram_size=0,
)
generation = generation[0]
if self.generation_includes_input:
generation = generation[input_len:]
decoded = self.processor.decode(
generation, skip_special_tokens=self.skip_special_tokens
)
return (decoded,)
def prepare_generation_params(
self, preprocessed_inputs: Dict[str, Any]
) -> Dict[str, Any]:
return preprocessed_inputs
def get_infer_bucket_file_list(self) -> list:
"""Get the list of required files for inference.
Returns:
list: A list of required files for inference, e.g., ["model.pt"].
"""
return [
"config.json",
"special_tokens_map.json",
"generation_config.json",
"tokenizer.json",
re.compile(r"model.*\.safetensors"),
"preprocessor_config.json",
"tokenizer_config.json",
]
def download_model_artifacts_from_roboflow_api(self) -> None:
api_data = get_roboflow_model_data(
api_key=self.api_key,
model_id=self.endpoint,
endpoint_type=ModelEndpointType.ORT,
device_id=self.device_id,
)
if "weights" not in api_data["ort"]:
raise ModelArtefactError(
f"`weights` key not available in Roboflow API response while downloading model weights."
)
for weights_url in api_data["ort"]["weights"].values():
t1 = perf_counter()
filename = weights_url.split("?")[0].split("/")[-1]
if filename.endswith(".npz"):
continue
model_weights_response = get_from_url(weights_url, json_response=False)
save_bytes_in_cache(
content=model_weights_response.content,
file=filename,
model_id=self.endpoint,
)
if filename.endswith("tar.gz"):
try:
subprocess.run(
[
"tar",
"-xzf",
os.path.join(self.cache_dir, filename),
"-C",
self.cache_dir,
],
check=True,
)
except subprocess.CalledProcessError as e:
raise ModelArtefactError(
f"Failed to extract model archive {filename}. Error: {str(e)}"
) from e
if perf_counter() - t1 > 120:
logger.debug(
"Weights download took longer than 120 seconds, refreshing API request"
)
api_data = get_roboflow_model_data(
api_key=self.api_key,
model_id=self.endpoint,
endpoint_type=ModelEndpointType.ORT,
device_id=self.device_id,
)
@property
def weights_file(self) -> None:
return None
def download_model_artefacts_from_s3(self) -> None:
raise NotImplementedError()
|