Bound video and audio masks to 8 bit level values to limit number of distinct token labels

This commit is contained in:
drozbay 2026-08-08 22:06:04 -06:00
parent 355e8d29d5
commit a2980c2927
1 changed files with 4 additions and 1 deletions

View File

@ -2202,11 +2202,14 @@ class MiniMaxH3(BaseModel):
video_mask = torch.nn.functional.pad(video_mask.reshape((-1,) + video_mask.shape[-3:]), (0, -w % pw, 0, -h % ph), mode="replicate")
video_mask = video_mask.reshape(lead + video_mask.shape[-2:])
video_mask = video_mask.reshape(video_mask.shape[:-2] + (video_mask.shape[-2] // ph, ph, video_mask.shape[-1] // pw, pw)).amax(dim=(-3, -1))
# values above 0.995 snap to 1.0, values below 0.05 snap to 0.0
# quantize to 1/256 so a gradient mask yields a bounded set of row timesteps
video_mask = torch.round(video_mask * 256.0) / 256.0
# threshold values above 0.995 to 1.0 and values below 0.05 to 0.0
video_mask = video_mask.masked_fill(video_mask >= 0.995, 1.0).masked_fill(video_mask <= 0.05, 0.0)
denoise_masks[0] = video_mask.repeat_interleave(ph, dim=-2).repeat_interleave(pw, dim=-1)[..., :h, :w]
if len(denoise_masks) > 1:
audio_mask = denoise_masks[1].amax(dim=1, keepdim=True)
audio_mask = torch.round(audio_mask * 256.0) / 256.0
audio_mask = audio_mask.masked_fill(audio_mask >= 0.995, 1.0).masked_fill(audio_mask <= 0.05, 0.0)
denoise_masks[1] = audio_mask.expand_as(denoise_masks[1]).contiguous()
return denoise_masks