diff --git a/videox_fun/data/dataset_image_video.py b/videox_fun/data/dataset_image_video.py index e1ef788..395cd05 100755 --- a/videox_fun/data/dataset_image_video.py +++ b/videox_fun/data/dataset_image_video.py @@ -759,7 +759,7 @@ class ImageVideoControlDataset(Dataset): if self.enable_inpaint and not self.enable_bucket: mask = get_random_mask(pixel_values.size()) - mask_pixel_values = pixel_values * (1 - mask) + torch.ones_like(pixel_values) * -1 * mask + mask_pixel_values = pixel_values * (1 - mask) + torch.zeros_like(pixel_values) * mask sample["mask_pixel_values"] = mask_pixel_values sample["mask"] = mask