Mask not fully zeroing out the background

[austinfujimori]

Hi guys, I'm applying a binary mask over the face of a person, using this code from the cli_lora_pti.py file:

    mask = (
        batch["mask"]
        .to(model_pred.device)
        .reshape(
            model_pred.shape[0], 1, model_pred.shape[2] * 8, model_pred.shape[3] * 8
        )
    )
    # resize to match model_pred
    mask = F.interpolate(
        mask.float(),
        size=model_pred.shape[-2:],
        mode="nearest",
    )

    mask = (mask + 0.01).pow(mask_temperature)

    mask = mask / mask.max()

    model_pred = model_pred * mask

    target = target * mask

loss = (
    F.mse_loss(model_pred.float(), target.float(), reduction="none")
    .mean([1, 2, 3])
    .mean()
). 

However, I notice during training that the background is still partially being fitted in. What could be causing this?

Some more information, all I do to the masks before this is convert it to 512x512 tensor, no mean/standard deviation normalization, pixel values from 0 to 255 are just scaled to 0 and 1.

The instance image pixel values however undergo these transformations before being encoded to 64x64 latent:

transforms.Resize(size, interpolation=transforms.InterpolationMode.BILINEAR),
transforms.CenterCrop(size) if center_crop else transforms.RandomCrop(size),
transforms.ToTensor(),
transforms.Normalize([0.5], [0.5]).

Thanks.