This commit is contained in:
Eugene Fairley 2025-06-27 12:19:33 -07:00
parent 73ce203ad9
commit 15c7750469

View File

@ -1089,20 +1089,20 @@ class WAN21(BaseModel):
image = image[:, :(extra_channels - 4)]
mask = kwargs.get("concat_mask", kwargs.get("denoise_mask", None))
if mask is None:
mask = torch.zeros_like(noise)[:, :4]
else:
if mask.shape[1] != 4:
mask = torch.mean(mask, dim=1, keepdim=True)
mask = 1.0 - mask
mask = utils.common_upscale(mask.to(device), noise.shape[-1], noise.shape[-2], "bilinear", "center")
if mask.shape[-3] < noise.shape[-3]:
mask = torch.nn.functional.pad(mask, (0, 0, 0, 0, 0, noise.shape[-3] - mask.shape[-3]), mode='constant', value=0)
if mask.shape[1] == 1:
mask = mask.repeat(1, 4, 1, 1, 1)
# if mask is None:
# mask = torch.zeros_like(noise)[:, :4]
# else:
# if mask.shape[1] != 4:
# mask = torch.mean(mask, dim=1, keepdim=True)
# mask = 1.0 - mask
# mask = utils.common_upscale(mask.to(device), noise.shape[-1], noise.shape[-2], "bilinear", "center")
# if mask.shape[-3] < noise.shape[-3]:
# mask = torch.nn.functional.pad(mask, (0, 0, 0, 0, 0, noise.shape[-3] - mask.shape[-3]), mode='constant', value=0)
# if mask.shape[1] == 1:
# mask = mask.repeat(1, 4, 1, 1, 1)
print(f"Mask shape: {mask.shape}, noise shape: {noise.shape}")
mask = utils.resize_to_batch_size(mask, noise.shape[0])
# print(f"Mask shape: {mask.shape}, noise shape: {noise.shape}")
# mask = utils.resize_to_batch_size(mask, noise.shape[0])
print(f"image shape: {image.shape}, mask shape: {mask.shape}")
return torch.cat((mask, image), dim=1)