| 192 | return (min_batch, max_batch, min_image_height, max_image_height, min_image_width, max_image_width, min_latent_height, max_latent_height, min_latent_width, max_latent_width) |
| 193 | |
| 194 | class CLIP(BaseModel): |
| 195 | def __init__(self, |
| 196 | hf_token, |
| 197 | device, |
| 198 | verbose, |
| 199 | path, |
| 200 | max_batch_size, |
| 201 | embedding_dim |
| 202 | ): |
| 203 | super(CLIP, self).__init__(hf_token, device=device, verbose=verbose, path=path, max_batch_size=max_batch_size, embedding_dim=embedding_dim) |
| 204 | self.name = "CLIP" |
| 205 | |
| 206 | def get_model(self): |
| 207 | return CLIPTextModel.from_pretrained(self.path, |
| 208 | subfolder="text_encoder", |
| 209 | use_auth_token=self.hf_token).to(self.device) |
| 210 | |
| 211 | def get_input_names(self): |
| 212 | return ['input_ids'] |
| 213 | |
| 214 | def get_output_names(self): |
| 215 | return ['text_embeddings', 'pooler_output'] |
| 216 | |
| 217 | def get_dynamic_axes(self): |
| 218 | return { |
| 219 | 'input_ids': {0: 'B'}, |
| 220 | 'text_embeddings': {0: 'B'} |
| 221 | } |
| 222 | |
| 223 | def get_input_profile(self, batch_size, image_height, image_width, static_batch, static_shape): |
| 224 | self.check_dims(batch_size, image_height, image_width) |
| 225 | min_batch, max_batch, _, _, _, _, _, _, _, _ = self.get_minmax_dims(batch_size, image_height, image_width, static_batch, static_shape) |
| 226 | return { |
| 227 | 'input_ids': [(min_batch, self.text_maxlen), (batch_size, self.text_maxlen), (max_batch, self.text_maxlen)] |
| 228 | } |
| 229 | |
| 230 | def get_shape_dict(self, batch_size, image_height, image_width): |
| 231 | self.check_dims(batch_size, image_height, image_width) |
| 232 | return { |
| 233 | 'input_ids': (batch_size, self.text_maxlen), |
| 234 | 'text_embeddings': (batch_size, self.text_maxlen, self.embedding_dim) |
| 235 | } |
| 236 | |
| 237 | def get_sample_input(self, batch_size, image_height, image_width): |
| 238 | self.check_dims(batch_size, image_height, image_width) |
| 239 | return torch.zeros(batch_size, self.text_maxlen, dtype=torch.int32, device=self.device) |
| 240 | |
| 241 | def optimize(self, onnx_graph): |
| 242 | opt = Optimizer(onnx_graph, verbose=self.verbose) |
| 243 | opt.info(self.name + ': original') |
| 244 | opt.select_outputs([0]) # delete graph output#1 |
| 245 | opt.cleanup() |
| 246 | opt.info(self.name + ': remove output[1]') |
| 247 | opt.fold_constants() |
| 248 | opt.info(self.name + ': fold constants') |
| 249 | opt.infer_shapes() |
| 250 | opt.info(self.name + ': shape inference') |
| 251 | opt.select_outputs([0], names=['text_embeddings']) # rename network output |