(self, inputs)
| 408 | |
| 409 | @torch.autocast('cuda', dtype=AUTOCAST_DTYPE) |
| 410 | def forward(self, inputs): |
| 411 | img, pe, txt_len, img_len, hw, vec_orig, *vec = inputs |
| 412 | |
| 413 | # Extract image portion (after text) |
| 414 | img = img[:, txt_len.item():, ...] |
| 415 | # Only process the noisy image portion (control image was concatenated after if present) |
| 416 | img = img[:, :img_len.item(), ...] |
| 417 | |
| 418 | out = self.final_layer(img, vec_orig) |
| 419 | out = out[:, :img_len.item()] |
| 420 | h_orig = hw[0].item() |
| 421 | w_orig = hw[1].item() |
| 422 | patch_size = self.patch_size |
| 423 | h_len = ((h_orig + (patch_size // 2)) // patch_size) |
| 424 | w_len = ((w_orig + (patch_size // 2)) // patch_size) |
| 425 | return rearrange(out, "b (h w) (c ph pw) -> b c (h ph) (w pw)", h=h_len, w=w_len, ph=patch_size, pw=patch_size)[:,:,:h_orig,:w_orig] |
nothing calls this directly
no outgoing calls
no test coverage detected