Face seg mask is not in [0, 1]

[nom]

Hi, I've been looking at the face segmentation masking feature. IIUC it should work such that only the face is being learnt and nothing else. So the loss of non-face pixels should be masked to 0.

However, when running the example use_face_conditioning_example.sh, and printing the mask values after

lora/lora_diffusion/cli_lora_pti.py

Lines 342 to 358 in bdd51b0

	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()

with
torch.unique(mask, return_counts=True), I'm seeing the lowest value is 0.35.
(tensor([0.3522, 0.3603, 0.4087, 0.4490, 1.0000], device='cuda:0'), tensor([6204, 1, 13, 13, 169], device='cuda:0'))

I see that the mask values are being adjusted here https://github.com/cloneofsimo/lora/blob/master/lora_diffusion/dataset.py#L288-L295
They're first being normalized to 0.5 mean, 0.5 std, and then multiplied by 0.5 and 1 is added, resulting in a 1.25 mean. Is this intended?