| 388 | return self.vae_encoder.encode(x).latent_dist.sample() |
| 389 | |
| 390 | class VAEEncoder(BaseModel): |
| 391 | def __init__(self, |
| 392 | hf_token, |
| 393 | device, |
| 394 | verbose, |
| 395 | path, |
| 396 | max_batch_size, |
| 397 | embedding_dim |
| 398 | ): |
| 399 | super(VAEEncoder, self).__init__(hf_token, device=device, verbose=verbose, path=path, max_batch_size=max_batch_size, embedding_dim=embedding_dim) |
| 400 | self.name = "VAE encoder" |
| 401 | |
| 402 | def get_model(self): |
| 403 | vae_encoder = TorchVAEEncoder(self.hf_token, self.device, self.path) |
| 404 | return vae_encoder |
| 405 | |
| 406 | def get_input_names(self): |
| 407 | return ['images'] |
| 408 | |
| 409 | def get_output_names(self): |
| 410 | return ['latent'] |
| 411 | |
| 412 | def get_dynamic_axes(self): |
| 413 | return { |
| 414 | 'images': {0: 'B', 2: '8H', 3: '8W'}, |
| 415 | 'latent': {0: 'B', 2: 'H', 3: 'W'} |
| 416 | } |
| 417 | |
| 418 | def get_input_profile(self, batch_size, image_height, image_width, static_batch, static_shape): |
| 419 | assert batch_size >= self.min_batch and batch_size <= self.max_batch |
| 420 | min_batch = batch_size if static_batch else self.min_batch |
| 421 | max_batch = batch_size if static_batch else self.max_batch |
| 422 | self.check_dims(batch_size, image_height, image_width) |
| 423 | min_batch, max_batch, min_image_height, max_image_height, min_image_width, max_image_width, _, _, _, _ = \ |
| 424 | self.get_minmax_dims(batch_size, image_height, image_width, static_batch, static_shape) |
| 425 | |
| 426 | return { |
| 427 | 'images': [(min_batch, 3, min_image_height, min_image_width), (batch_size, 3, image_height, image_width), (max_batch, 3, max_image_height, max_image_width)], |
| 428 | } |
| 429 | |
| 430 | def get_shape_dict(self, batch_size, image_height, image_width): |
| 431 | latent_height, latent_width = self.check_dims(batch_size, image_height, image_width) |
| 432 | return { |
| 433 | 'images': (batch_size, 3, image_height, image_width), |
| 434 | 'latent': (batch_size, 4, latent_height, latent_width) |
| 435 | } |
| 436 | |
| 437 | def get_sample_input(self, batch_size, image_height, image_width): |
| 438 | self.check_dims(batch_size, image_height, image_width) |
| 439 | return torch.randn(batch_size, 3, image_height, image_width, dtype=torch.float32, device=self.device) |
| 440 | |
| 441 | def make_VAEEncoder(version, hf_token, device, verbose, max_batch_size, inpaint=False): |
| 442 | return VAEEncoder(hf_token=hf_token, device=device, verbose=verbose, path=get_path(version, inpaint=inpaint), |