(self, sd_vae, latent)
| 381 | |
| 382 | @torch.no_grad() |
| 383 | def forward(self, sd_vae, latent): |
| 384 | pixel = sd_vae.decode(latent).sample |
| 385 | pixel = (pixel * 0.5 + 0.5).clip(0, 1).to(self.dtype) |
| 386 | latent = latent.to(self.dtype) |
| 387 | result_list = [] |
| 388 | vis_list = [] |
| 389 | |
| 390 | for i in range(int(latent.shape[0])): |
| 391 | y = self.estimate_augmented(pixel[i:i + 1], latent[i:i + 1]) |
| 392 | |
| 393 | y = y.clip(0, 1).movedim(1, -1) |
| 394 | alpha = y[..., :1] |
| 395 | fg = y[..., 1:] |
| 396 | |
| 397 | B, H, W, C = fg.shape |
| 398 | cb = checkerboard(shape=(H // 64, W // 64)) |
| 399 | cb = cv2.resize(cb, (W, H), interpolation=cv2.INTER_NEAREST) |
| 400 | cb = (0.5 + (cb - 0.5) * 0.1)[None, ..., None] |
| 401 | cb = torch.from_numpy(cb).to(fg) |
| 402 | |
| 403 | vis = (fg * alpha + cb * (1 - alpha))[0] |
| 404 | vis = (vis * 255.0).detach().cpu().float().numpy().clip(0, 255).astype(np.uint8) |
| 405 | vis_list.append(vis) |
| 406 | |
| 407 | png = torch.cat([fg, alpha], dim=3)[0] |
| 408 | png = (png * 255.0).detach().cpu().float().numpy().clip(0, 255).astype(np.uint8) |
| 409 | result_list.append(png) |
| 410 | |
| 411 | return result_list, vis_list |
| 412 | |
| 413 | |
| 414 | class TransparentVAEEncoder(torch.nn.Module): |
nothing calls this directly
no test coverage detected