| 108 | |
| 109 | |
| 110 | class FastTileWorker: |
| 111 | def __init__(self): |
| 112 | pass |
| 113 | |
| 114 | |
| 115 | def build_mask(self, data, is_bound): |
| 116 | _, _, H, W = data.shape |
| 117 | h = repeat(torch.arange(H), "H -> H W", H=H, W=W) |
| 118 | w = repeat(torch.arange(W), "W -> H W", H=H, W=W) |
| 119 | border_width = (H + W) // 4 |
| 120 | pad = torch.ones_like(h) * border_width |
| 121 | mask = torch.stack([ |
| 122 | pad if is_bound[0] else h + 1, |
| 123 | pad if is_bound[1] else H - h, |
| 124 | pad if is_bound[2] else w + 1, |
| 125 | pad if is_bound[3] else W - w |
| 126 | ]).min(dim=0).values |
| 127 | mask = mask.clip(1, border_width) |
| 128 | mask = (mask / border_width).to(dtype=data.dtype, device=data.device) |
| 129 | mask = rearrange(mask, "H W -> 1 H W") |
| 130 | return mask |
| 131 | |
| 132 | |
| 133 | def tiled_forward(self, forward_fn, model_input, tile_size, tile_stride, tile_device="cpu", tile_dtype=torch.float32, border_width=None): |
| 134 | # Prepare |
| 135 | B, C, H, W = model_input.shape |
| 136 | border_width = int(tile_stride*0.5) if border_width is None else border_width |
| 137 | weight = torch.zeros((1, 1, H, W), dtype=tile_dtype, device=tile_device) |
| 138 | values = torch.zeros((B, C, H, W), dtype=tile_dtype, device=tile_device) |
| 139 | |
| 140 | # Split tasks |
| 141 | tasks = [] |
| 142 | for h in range(0, H, tile_stride): |
| 143 | for w in range(0, W, tile_stride): |
| 144 | if (h-tile_stride >= 0 and h-tile_stride+tile_size >= H) or (w-tile_stride >= 0 and w-tile_stride+tile_size >= W): |
| 145 | continue |
| 146 | h_, w_ = h + tile_size, w + tile_size |
| 147 | if h_ > H: h, h_ = H - tile_size, H |
| 148 | if w_ > W: w, w_ = W - tile_size, W |
| 149 | tasks.append((h, h_, w, w_)) |
| 150 | |
| 151 | # Run |
| 152 | for hl, hr, wl, wr in tasks: |
| 153 | # Forward |
| 154 | hidden_states_batch = forward_fn(hl, hr, wl, wr).to(dtype=tile_dtype, device=tile_device) |
| 155 | |
| 156 | mask = self.build_mask(hidden_states_batch, is_bound=(hl==0, hr>=H, wl==0, wr>=W)) |
| 157 | values[:, :, hl:hr, wl:wr] += hidden_states_batch * mask |
| 158 | weight[:, :, hl:hr, wl:wr] += mask |
| 159 | values /= weight |
| 160 | return values |
| 161 | |
| 162 | |
| 163 | |