-
Notifications
You must be signed in to change notification settings - Fork 4
/
tiling.py
180 lines (151 loc) · 7.4 KB
/
tiling.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
import torch
import itertools
import numpy as np
# tiled sampler code adapted from https://github.com/BlenderNeko/ComfyUI_TiledKSampler
# for use with https://github.com/ClownsharkBatwing/UltraCascade
def grouper(n, iterable):
it = iter(iterable)
while True:
chunk = list(itertools.islice(it, n))
if not chunk:
return
yield chunk
def create_batches(n, iterable):
groups = itertools.groupby(iterable, key= lambda x: (x[1], x[3]))
for _, x in groups:
for y in grouper(n, x):
yield y
def get_slice(tensor, h, h_len, w, w_len):
t = tensor.narrow(-2, h, h_len)
t = t.narrow(-1, w, w_len)
return t
def set_slice(tensor1,tensor2, h, h_len, w, w_len, mask=None):
if mask is not None:
tensor1[:,:,h:h+h_len,w:w+w_len] = tensor1[:,:,h:h+h_len,w:w+w_len] * (1 - mask) + tensor2 * mask
else:
tensor1[:,:,h:h+h_len,w:w+w_len] = tensor2
def get_tiles_and_masks_simple(steps, latent_shape, tile_height, tile_width, compression=4):
latent_size_h = latent_shape[-2]
latent_size_w = latent_shape[-1]
tile_size_h = int(tile_height // compression) #CHANGED FROM 8
tile_size_w = int(tile_width // compression) #CHANGED FROM 8
h = np.arange(0,latent_size_h, tile_size_h)
w = np.arange(0,latent_size_w, tile_size_w)
def create_tile(hs, ws, i, j):
h = int(hs[i])
w = int(ws[j])
h_len = min(tile_size_h, latent_size_h - h)
w_len = min(tile_size_w, latent_size_w - w)
return (h, h_len, w, w_len, steps, None)
passes = [
[[create_tile(h, w, i, j) for i in range(len(h)) for j in range(len(w))]],
]
return passes
def get_tiles_and_masks_padded(steps, latent_shape, tile_height, tile_width, compression=4):
batch_size = latent_shape[0]
latent_size_h = latent_shape[-2]
latent_size_w = latent_shape[-1]
tile_size_h = int(tile_height // compression) #CHANGED FROM 8
tile_size_w = int(tile_width // compression) #CHANGED FROM 8
#if compression > 1:
tile_size_h = int((tile_size_h // 4) * 4) #MIGHT BE A PROBLEM WITH STAGE C?
tile_size_w = int((tile_size_w // 4) * 4)
#masks
mask_h = [0,tile_size_h // 4, tile_size_h - tile_size_h // 4, tile_size_h]
mask_w = [0,tile_size_w // 4, tile_size_w - tile_size_w // 4, tile_size_w]
masks = [[] for _ in range(3)]
for i in range(3):
for j in range(3):
mask = torch.zeros((batch_size,1,tile_size_h, tile_size_w), dtype=torch.float32, device='cpu')
mask[:,:, mask_h[i]:mask_h[i+1],
mask_w[j]:mask_w[j+1]] = 1.0
masks[i].append(mask)
def create_mask(h_ind, w_ind, h_ind_max, w_ind_max, mask_h, mask_w, h_len, w_len):
mask = masks[1][1]
if not (h_ind == 0 or h_ind == h_ind_max or w_ind == 0 or w_ind == w_ind_max):
return get_slice(mask, 0, h_len, 0, w_len)
mask = mask.clone()
if h_ind == 0 and mask_h:
mask += masks[0][1]
if h_ind == h_ind_max and mask_h:
mask += masks[2][1]
if w_ind == 0 and mask_w:
mask += masks[1][0]
if w_ind == w_ind_max and mask_w:
mask += masks[1][2]
if h_ind == 0 and w_ind == 0 and mask_h and mask_w:
mask += masks[0][0]
if h_ind == 0 and w_ind == w_ind_max and mask_h and mask_w:
mask += masks[0][2]
if h_ind == h_ind_max and w_ind == 0 and mask_h and mask_w:
mask += masks[2][0]
if h_ind == h_ind_max and w_ind == w_ind_max and mask_h and mask_w:
mask += masks[2][2]
return get_slice(mask, 0, h_len, 0, w_len)
h = np.arange(0,latent_size_h, tile_size_h)
h_shift = np.arange(tile_size_h // 2, latent_size_h - tile_size_h // 2, tile_size_h)
w = np.arange(0,latent_size_w, tile_size_w)
w_shift = np.arange(tile_size_w // 2, latent_size_w - tile_size_h // 2, tile_size_w)
def create_tile(hs, ws, mask_h, mask_w, i, j):
h = int(hs[i])
w = int(ws[j])
h_len = min(tile_size_h, latent_size_h - h)
w_len = min(tile_size_w, latent_size_w - w)
mask = create_mask(i,j,len(hs)-1, len(ws)-1, mask_h, mask_w, h_len, w_len)
return (h, h_len, w, w_len, steps, mask)
passes = [
[[create_tile(h, w, True, True, i, j) for i in range(len(h)) for j in range(len(w))]],
[[create_tile(h_shift, w, False, True, i, j) for i in range(len(h_shift)) for j in range(len(w))]],
[[create_tile(h, w_shift, True, False, i, j) for i in range(len(h)) for j in range(len(w_shift))]],
[[create_tile(h_shift, w_shift, False, False, i,j) for i in range(len(h_shift)) for j in range(len(w_shift))]],
]
return passes
def mask_at_boundary(h, h_len, w, w_len, tile_size_h, tile_size_w, latent_size_h, latent_size_w, mask, device='cpu', compression=4):
tile_size_h = int(tile_size_h // compression) #CHANGED FROM 8
tile_size_w = int(tile_size_w // compression) #CHANGED FROM 8
if (h_len == tile_size_h or h_len == latent_size_h) and (w_len == tile_size_w or w_len == latent_size_w):
return h, h_len, w, w_len, mask
h_offset = min(0, latent_size_h - (h + tile_size_h))
w_offset = min(0, latent_size_w - (w + tile_size_w))
new_mask = torch.zeros((1,1,tile_size_h, tile_size_w), dtype=torch.float32, device=device)
new_mask[:,:,-h_offset:h_len if h_offset == 0 else tile_size_h, -w_offset:w_len if w_offset == 0 else tile_size_w] = 1.0 if mask is None else mask
return h + h_offset, tile_size_h, w + w_offset, tile_size_w, new_mask
def get_tiles_and_masks_rgrid(steps, latent_shape, tile_height, tile_width, generator, compression=4):
def calc_coords(latent_size, tile_size, jitter):
tile_coords = int((latent_size + jitter - 1) // tile_size + 1)
tile_coords = [np.clip(tile_size * c - jitter, 0, latent_size) for c in range(tile_coords + 1)]
tile_coords = [(c1, c2-c1) for c1, c2 in zip(tile_coords, tile_coords[1:])]
return tile_coords
#calc stuff
batch_size = latent_shape[0]
latent_size_h = latent_shape[-2]
latent_size_w = latent_shape[-1]
tile_size_h = int(tile_height // compression) #CHANGED FROM 8
tile_size_w = int(tile_width // compression) #CHANGED FROM 8
tiles_all = []
for s in range(steps):
rands = torch.rand((2,), dtype=torch.float32, generator=generator, device='cpu').numpy()
jitter_w1 = int(rands[0] * tile_size_w)
jitter_w2 = int(((rands[0] + .5) % 1.0) * tile_size_w)
jitter_h1 = int(rands[1] * tile_size_h)
jitter_h2 = int(((rands[1] + .5) % 1.0) * tile_size_h)
#calc number of tiles
tiles_h = [
calc_coords(latent_size_h, tile_size_h, jitter_h1),
calc_coords(latent_size_h, tile_size_h, jitter_h2)
]
tiles_w = [
calc_coords(latent_size_w, tile_size_w, jitter_w1),
calc_coords(latent_size_w, tile_size_w, jitter_w2)
]
tiles = []
if s % 2 == 0:
for i, h in enumerate(tiles_h[0]):
for w in tiles_w[i%2]:
tiles.append((int(h[0]), int(h[1]), int(w[0]), int(w[1]), 1, None))
else:
for i, w in enumerate(tiles_w[0]):
for h in tiles_h[i%2]:
tiles.append((int(h[0]), int(h[1]), int(w[0]), int(w[1]), 1, None))
tiles_all.append(tiles)
return [tiles_all]