update TileWorker for better visual quality

diffsynth/models/tiler.py

@@ -2,76 +2,6 @@ import torch
 from einops import rearrange, repeat
 
 
-class Tiler(torch.nn.Module):
-    def __init__(self):
-        super().__init__()
-
-    def mask(self, height, width, line_width):
-        x = torch.arange(height).repeat(width, 1).T
-        y = torch.arange(width).repeat(height, 1)
-        mask = torch.stack([x + 1, height - x, y + 1, width - y]).min(dim=0).values
-        mask = (mask / line_width).clip(0, 1)
-        return mask
-
-    def forward(self, forward_fn, x, tile_size, tile_stride, batch_size=1, inter_device="cpu", inter_dtype=torch.float32):
-        # Prepare
-        device = x.device
-        torch_dtype = x.dtype
-
-        # tile
-        b, c_in, h_in, w_in = x.shape
-        x = x.to(device=inter_device, dtype=inter_dtype)
-        fold_params = {
-            "kernel_size": (tile_size, tile_size),
-            "stride": (tile_stride, tile_stride)
-        }
-        unfold_operator = torch.nn.Unfold(**fold_params)
-        x = unfold_operator(x)
-        x = x.view((b, c_in, tile_size, tile_size, -1))
-
-        # inference
-        x_out_stack = []
-        for tile_id in range(0, x.shape[-1], batch_size):
-
-            # process input
-            next_tile_id = min(tile_id + batch_size, x.shape[-1])
-            x_in = x[:, :, :, :, tile_id: next_tile_id]
-            x_in = x_in.to(device=device, dtype=torch_dtype)
-            x_in = x_in.permute(4, 0, 1, 2, 3)
-            x_in = x_in.view((x_in.shape[0]*x_in.shape[1], x_in.shape[2], x_in.shape[3], x_in.shape[4]))
-
-            # process output
-            x_out = forward_fn(x_in)
-            x_out = x_out.view((next_tile_id - tile_id, b, x_out.shape[1], x_out.shape[2], x_out.shape[3]))
-            x_out = x_out.permute(1, 2, 3, 4, 0)
-            x_out = x_out.to(device=inter_device, dtype=inter_dtype)
-            x_out_stack.append(x_out)
-
-        x = torch.concat(x_out_stack, dim=-1)
-
-        # untile
-        in2out_scale = x.shape[2] / tile_size
-        h_out, w_out = int(h_in * in2out_scale), int(w_in * in2out_scale)
-
-        mask = self.mask(int(tile_size * in2out_scale), int(tile_size * in2out_scale), int(tile_stride * in2out_scale * 0.5))
-        mask = mask.to(device=inter_device, dtype=inter_dtype)
-        mask = mask.reshape((1, 1, mask.shape[0], mask.shape[1], 1))
-        x = x * mask
-
-        fold_params = {
-            "kernel_size": (int(tile_size * in2out_scale), int(tile_size * in2out_scale)),
-            "stride": (int(tile_stride * in2out_scale), int(tile_stride * in2out_scale))
-        }
-        fold_operator = torch.nn.Fold(output_size=(h_out, w_out), **fold_params)
-        divisor = fold_operator(mask.repeat(1, 1, 1, 1, x.shape[-1]).view(b, -1, x.shape[-1]))
-
-        x = x.view((b, -1, x.shape[-1]))
-        x = fold_operator(x) / divisor
-        x = x.to(device=device, dtype=torch_dtype)
-
-        return x
-
-
 class TileWorker:
     def __init__(self):
         pass