(self, sd_vae, list_of_np_rgba_hwc_uint8, use_offset=True)
| 430 | |
| 431 | @torch.no_grad() |
| 432 | def forward(self, sd_vae, list_of_np_rgba_hwc_uint8, use_offset=True): |
| 433 | list_of_np_rgb_padded = [pad_rgb(x) for x in list_of_np_rgba_hwc_uint8] |
| 434 | rgb_padded_bchw_01 = torch.from_numpy(np.stack(list_of_np_rgb_padded, axis=0)).float().movedim(-1, 1) |
| 435 | rgba_bchw_01 = torch.from_numpy(np.stack(list_of_np_rgba_hwc_uint8, axis=0)).float().movedim(-1, 1) / 255.0 |
| 436 | rgb_bchw_01 = rgba_bchw_01[:, :3, :, :] |
| 437 | a_bchw_01 = rgba_bchw_01[:, 3:, :, :] |
| 438 | vae_feed = (rgb_bchw_01 * 2.0 - 1.0) * a_bchw_01 |
| 439 | vae_feed = vae_feed.to(device=sd_vae.device, dtype=sd_vae.dtype) |
| 440 | latent_dist = sd_vae.encode(vae_feed).latent_dist |
| 441 | offset_feed = torch.cat([a_bchw_01, rgb_padded_bchw_01], dim=1).to(device=sd_vae.device, dtype=self.dtype) |
| 442 | offset = self.model(offset_feed) * self.alpha |
| 443 | if use_offset: |
| 444 | latent = dist_sample_deterministic(dist=latent_dist, perturbation=offset) |
| 445 | else: |
| 446 | latent = latent_dist.sample() |
| 447 | return latent |
nothing calls this directly
no test coverage detected