There is a difference between training code and test code

[Afakash]

Describe the bug

Hello, I found that there may be a bug in the BrushNet training code in the example.
[code file: /examples/brushnet/train_brushnet.py ]

Specifically, the white area (value = 1) of the mask generated from the datasets is align with the unmasked area (that is, encoded by VAE into conditional latent as the input of BrushNet branch).

In the test code, the unmasked area is aligned with the black area in the mask (Value=0). Especially in the test code, the white area corresponds to the part that needs to be filled.
[code file: /examples/brushnet/test_brushnet.py ]

In addition, in your paper, the black area is inpainted in Figure 1, while in Figure 3, the white area of ​​the mask is inpainted.

Reproduction

// train code:
masks = torch.nn.functional.interpolate(
batch["masks"],
size=(
latents.shape[-2],
latents.shape[-1]
)
)
// test code:
init_image = init_image * (1-mask_image)

Logs

No response

System Info

cuda: 11.8
torch: 2.3.0

Who can help?

No response