(self, model_output, height, width, tile_size, tile_stride, border_width, tile_device, tile_dtype)
| 62 | |
| 63 | |
| 64 | def untile(self, model_output, height, width, tile_size, tile_stride, border_width, tile_device, tile_dtype): |
| 65 | # The reversed function of tile |
| 66 | mask = self.mask(tile_size, tile_size, border_width) |
| 67 | mask = mask.to(device=tile_device, dtype=tile_dtype) |
| 68 | mask = rearrange(mask, "h w -> 1 1 h w 1") |
| 69 | model_output = model_output * mask |
| 70 | |
| 71 | fold_operator = torch.nn.Fold( |
| 72 | output_size=(height, width), |
| 73 | kernel_size=(tile_size, tile_size), |
| 74 | stride=(tile_stride, tile_stride) |
| 75 | ) |
| 76 | mask = repeat(mask[0, 0, :, :, 0], "h w -> 1 (h w) n", n=model_output.shape[-1]) |
| 77 | model_output = rearrange(model_output, "b c h w n -> b (c h w) n") |
| 78 | model_output = fold_operator(model_output) / fold_operator(mask) |
| 79 | |
| 80 | return model_output |
| 81 | |
| 82 | |
| 83 | def tiled_forward(self, forward_fn, model_input, tile_size, tile_stride, tile_batch_size=1, tile_device="cpu", tile_dtype=torch.float32, border_width=None): |
no test coverage detected