fix custom mask in tiled batch

This commit is contained in:
matt3o
2024-03-27 09:28:58 +01:00
parent 477217b364
commit 83bf9108ba
+17 -20
View File
@@ -672,14 +672,16 @@ class IPAdapterTiled:
del ipadapter
# 2. Extract the tiles
tile_size = 256 # I'm using 256 instead of 224 as it is more likely divisible by the latent size, it will be downscaled to 224 by the clip vision encoder
_, oh, ow, _ = image.shape
if attn_mask is None:
attn_mask = torch.ones([1, oh, ow], dtype=image.dtype, device=image.device)
# 2. Extract the tiles
image = image.permute([0,3,1,2])
if attn_mask is not None:
attn_mask = attn_mask.unsqueeze(1)
# the mask should have the same proportions as the reference image and the latent
attn_mask = T.Resize((oh, ow), interpolation=T.InterpolationMode.BICUBIC, antialias=True)(attn_mask)
attn_mask = attn_mask.unsqueeze(1)
# the mask should have the same proportions as the reference image and the latent
attn_mask = T.Resize((oh, ow), interpolation=T.InterpolationMode.BICUBIC, antialias=True)(attn_mask)
# if the image is almost a square, we crop it to a square
if oh / ow > 0.75 and oh / ow < 1.33:
@@ -687,32 +689,30 @@ class IPAdapterTiled:
image = T.CenterCrop(min(oh, ow))(image)
resize = (tile_size*2, tile_size*2)
if attn_mask is not None:
attn_mask = T.CenterCrop(min(oh, ow))(attn_mask)
# otherwise resize the smallest side to 224 and the other proportionally
attn_mask = T.CenterCrop(min(oh, ow))(attn_mask)
# otherwise resize the smallest side and the other proportionally
else:
resize = (int(tile_size * ow / oh), tile_size) if oh < ow else (tile_size, int(tile_size * oh / ow))
# using PIL for better results
imgs = []
for img in image:
img = T.ToPILImage()(img) # using PIL for better results
img = T.ToPILImage()(img)
img = img.resize(resize, resample=Image.Resampling['LANCZOS'])
imgs.append(T.ToTensor()(img))
image = torch.stack(imgs)
del imgs, img
if attn_mask is not None:
attn_mask = T.Resize(resize[::-1], interpolation=T.InterpolationMode.BICUBIC, antialias=True)(attn_mask) # we don't need a high quality resize for the mask
# we don't need a high quality resize for the mask
attn_mask = T.Resize(resize[::-1], interpolation=T.InterpolationMode.BICUBIC, antialias=True)(attn_mask)
# we allow a maximum of 4 tiles
if oh / ow > 4 or oh / ow < 0.25:
crop = (tile_size, tile_size*4) if oh < ow else (tile_size*4, tile_size)
image = T.CenterCrop(crop)(image)
if attn_mask is not None:
attn_mask = T.CenterCrop(crop)(attn_mask)
attn_mask = T.CenterCrop(crop)(attn_mask)
if attn_mask is not None:
attn_mask = attn_mask.squeeze(1)
attn_mask = attn_mask.squeeze(1)
if sharpening > 0:
image = contrast_adaptive_sharpening(image, sharpening)
@@ -727,7 +727,7 @@ class IPAdapterTiled:
overlap_x = max(0, (tiles_x * tile_size - ow) / (tiles_x - 1 if tiles_x > 1 else 1))
overlap_y = max(0, (tiles_y * tile_size - oh) / (tiles_y - 1 if tiles_y > 1 else 1))
base_mask = torch.zeros([1, oh, ow], dtype=image.dtype, device=image.device)
base_mask = torch.zeros([attn_mask.shape[0], oh, ow], dtype=image.dtype, device=image.device)
# extract all the tiles from the image and create the masks
tiles = []
@@ -738,10 +738,7 @@ class IPAdapterTiled:
start_y = int(y * (tile_size - overlap_y))
tiles.append(image[:, start_y:start_y+tile_size, start_x:start_x+tile_size, :])
mask = base_mask.clone()
if attn_mask is not None:
mask[:, start_y:start_y+tile_size, start_x:start_x+tile_size] = attn_mask[:, start_y:start_y+tile_size, start_x:start_x+tile_size]
else:
mask[:, start_y:start_y+tile_size, start_x:start_x+tile_size] = 1
mask[:, start_y:start_y+tile_size, start_x:start_x+tile_size] = attn_mask[:, start_y:start_y+tile_size, start_x:start_x+tile_size]
masks.append(mask)
del mask