|
|
|
|
@ -23,8 +23,19 @@ def sampling_function(model_function, x, timestep, uncond, cond, cond_scale, con
|
|
|
|
|
adm_cond = cond[1]['adm_encoded'] |
|
|
|
|
|
|
|
|
|
input_x = x_in[:,:,area[2]:area[0] + area[2],area[3]:area[1] + area[3]] |
|
|
|
|
mult = torch.ones_like(input_x) * strength |
|
|
|
|
if 'mask' in cond[1]: |
|
|
|
|
# Scale the mask to the size of the input |
|
|
|
|
# The mask should have been resized as we began the sampling process |
|
|
|
|
mask = cond[1]['mask'] |
|
|
|
|
assert(mask.shape[1] == x_in.shape[2]) |
|
|
|
|
assert(mask.shape[2] == x_in.shape[3]) |
|
|
|
|
mask = mask[:,area[2]:area[0] + area[2],area[3]:area[1] + area[3]] |
|
|
|
|
mask = mask.unsqueeze(1).repeat(input_x.shape[0] // mask.shape[0], input_x.shape[1], 1, 1) |
|
|
|
|
else: |
|
|
|
|
mask = torch.ones_like(input_x) |
|
|
|
|
mult = mask * strength |
|
|
|
|
|
|
|
|
|
if 'mask' not in cond[1]: |
|
|
|
|
rr = 8 |
|
|
|
|
if area[2] != 0: |
|
|
|
|
for t in range(rr): |
|
|
|
|
@ -38,6 +49,7 @@ def sampling_function(model_function, x, timestep, uncond, cond, cond_scale, con
|
|
|
|
|
if (area[1] + area[3]) < x_in.shape[3]: |
|
|
|
|
for t in range(rr): |
|
|
|
|
mult[:,:,:,area[1] - 1 - t:area[1] - t] *= ((1.0/rr) * (t + 1)) |
|
|
|
|
|
|
|
|
|
conditionning = {} |
|
|
|
|
conditionning['c_crossattn'] = cond[0] |
|
|
|
|
if cond_concat_in is not None and len(cond_concat_in) > 0: |
|
|
|
|
@ -301,6 +313,71 @@ def blank_inpaint_image_like(latent_image):
|
|
|
|
|
blank_image[:,3] *= 0.1380 |
|
|
|
|
return blank_image |
|
|
|
|
|
|
|
|
|
def get_mask_aabb(masks): |
|
|
|
|
if masks.numel() == 0: |
|
|
|
|
return torch.zeros((0, 4), device=masks.device, dtype=torch.int) |
|
|
|
|
|
|
|
|
|
b = masks.shape[0] |
|
|
|
|
|
|
|
|
|
bounding_boxes = torch.zeros((b, 4), device=masks.device, dtype=torch.int) |
|
|
|
|
is_empty = torch.zeros((b), device=masks.device, dtype=torch.bool) |
|
|
|
|
for i in range(b): |
|
|
|
|
mask = masks[i] |
|
|
|
|
if mask.numel() == 0: |
|
|
|
|
continue |
|
|
|
|
if torch.max(mask != 0) == False: |
|
|
|
|
is_empty[i] = True |
|
|
|
|
continue |
|
|
|
|
y, x = torch.where(mask) |
|
|
|
|
bounding_boxes[i, 0] = torch.min(x) |
|
|
|
|
bounding_boxes[i, 1] = torch.min(y) |
|
|
|
|
bounding_boxes[i, 2] = torch.max(x) |
|
|
|
|
bounding_boxes[i, 3] = torch.max(y) |
|
|
|
|
|
|
|
|
|
return bounding_boxes, is_empty |
|
|
|
|
|
|
|
|
|
def resolve_cond_masks(conditions, h, w, device): |
|
|
|
|
# We need to decide on an area outside the sampling loop in order to properly generate opposite areas of equal sizes. |
|
|
|
|
# While we're doing this, we can also resolve the mask device and scaling for performance reasons |
|
|
|
|
for i in range(len(conditions)): |
|
|
|
|
c = conditions[i] |
|
|
|
|
if 'mask' in c[1]: |
|
|
|
|
mask = c[1]['mask'] |
|
|
|
|
mask = mask.to(device=device) |
|
|
|
|
modified = c[1].copy() |
|
|
|
|
if len(mask.shape) == 2: |
|
|
|
|
mask = mask.unsqueeze(0) |
|
|
|
|
if mask.shape[2] != h or mask.shape[3] != w: |
|
|
|
|
mask = torch.nn.functional.interpolate(mask.unsqueeze(1), size=(h, w), mode='bilinear', align_corners=False).squeeze(1) |
|
|
|
|
|
|
|
|
|
if modified.get("set_area_to_bounds", False): |
|
|
|
|
bounds = torch.max(torch.abs(mask),dim=0).values.unsqueeze(0) |
|
|
|
|
boxes, is_empty = get_mask_aabb(bounds) |
|
|
|
|
if is_empty[0]: |
|
|
|
|
# Use the minimum possible size for efficiency reasons. (Since the mask is all-0, this becomes a noop anyway) |
|
|
|
|
modified['area'] = (8, 8, 0, 0) |
|
|
|
|
else: |
|
|
|
|
box = boxes[0] |
|
|
|
|
H, W, Y, X = (box[3] - box[1] + 1, box[2] - box[0] + 1, box[1], box[0]) |
|
|
|
|
# Make sure the height and width are divisible by 8 |
|
|
|
|
if X % 8 != 0: |
|
|
|
|
newx = X // 8 * 8 |
|
|
|
|
W = W + (X - newx) |
|
|
|
|
X = newx |
|
|
|
|
if Y % 8 != 0: |
|
|
|
|
newy = Y // 8 * 8 |
|
|
|
|
H = H + (Y - newy) |
|
|
|
|
Y = newy |
|
|
|
|
if H % 8 != 0: |
|
|
|
|
H = H + (8 - (H % 8)) |
|
|
|
|
if W % 8 != 0: |
|
|
|
|
W = W + (8 - (W % 8)) |
|
|
|
|
area = (int(H), int(W), int(Y), int(X)) |
|
|
|
|
modified['area'] = area |
|
|
|
|
|
|
|
|
|
modified['mask'] = mask |
|
|
|
|
conditions[i] = [c[0], modified] |
|
|
|
|
|
|
|
|
|
def create_cond_with_same_area_if_none(conds, c): |
|
|
|
|
if 'area' not in c[1]: |
|
|
|
|
return |
|
|
|
|
@ -461,7 +538,6 @@ class KSampler:
|
|
|
|
|
sigmas = self.calculate_sigmas(new_steps).to(self.device) |
|
|
|
|
self.sigmas = sigmas[-(steps + 1):] |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def sample(self, noise, positive, negative, cfg, latent_image=None, start_step=None, last_step=None, force_full_denoise=False, denoise_mask=None, sigmas=None, callback=None): |
|
|
|
|
if sigmas is None: |
|
|
|
|
sigmas = self.sigmas |
|
|
|
|
@ -484,6 +560,10 @@ class KSampler:
|
|
|
|
|
|
|
|
|
|
positive = positive[:] |
|
|
|
|
negative = negative[:] |
|
|
|
|
|
|
|
|
|
resolve_cond_masks(positive, noise.shape[2], noise.shape[3], self.device) |
|
|
|
|
resolve_cond_masks(negative, noise.shape[2], noise.shape[3], self.device) |
|
|
|
|
|
|
|
|
|
#make sure each cond area has an opposite one with the same area |
|
|
|
|
for c in positive: |
|
|
|
|
create_cond_with_same_area_if_none(negative, c) |
|
|
|
|
|
comfyanonymous
2 years ago