RomanNum3ral
/
ComfyUI
mirror of https://github.com/Comfy-Org/ComfyUI.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Allow fractional mask values for MiniMax-H3

This commit is contained in:
drozbay 2026-08-08 15:14:42 -06:00
parent 877b5b8862
commit 28aa924e81
2 changed files with 57 additions and 46 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

73
comfy/ldm/minimax/model.py
View File

@ -74,14 +74,15 @@ def _axis_from_sqrt_area(dim, patch, sqrt_area):
return (torch.arange(n, dtype=torch.float64) * (ratio / n) + (1.0 - ratio) / 2.0) * 32.0
def mask_row_targets(mask, latent_t, lat_h, lat_w):
# [T, H, W] denoise mask (1 = generate) -> per-2x2-patch-row bool, None when every row generates
def mask_row_values(mask, latent_t, lat_h, lat_w):
# [T, H, W] denoise mask (1 = generate) -> per-2x2-patch-row float in [0, 1],
# None when every row fully generates
m = torch.nn.functional.pad(mask, (0, lat_w - mask.shape[-1], 0, lat_h - mask.shape[-2]), mode="replicate")
m = m.reshape(latent_t, lat_h // 2, 2, lat_w // 2, 2).amax(dim=(2, 4))
target = m.reshape(-1) >= 0.5
if bool(target.all()):
values = m.reshape(-1)
if bool((values >= 1.0 - 1e-3).all()):
return None
return target
return values
def _frame_grid(h, w):
@ -210,10 +211,7 @@ class AdalnProj(nn.Module):
def _mod_row(vecs, row, dtype):
# row is a mod-row index, or (target_row, pin_row, weight[n,1]) blending two rows per token
if isinstance(row, tuple):
rt, rp, w = row
return torch.lerp(vecs[rp], vecs[rt], w.to(vecs.dtype)).to(dtype)
# row is a mod-row index, or a per-token LongTensor of mod-row indices
return vecs[row].to(dtype)
@ -566,32 +564,43 @@ class MiniMaxH3Model(nn.Module):
"cond": max(t_v, vis_aug), "ref_img": max(t_v, vis_aug),
"ref_audio": max(t_a, aud_aug)}
# rows that are preserved by the noise mask run at the cond timestep
# masked rows run at their own strength: mask value m puts a row at sigma = m * sigma_stream,
# so its label is 1 - m * sigma, clamped at the cond timestep for fully preserved rows
t_pin_v = max(t_v, VISUAL_COND_TIMESTEP)
t_pin_a = max(t_a, AUDIO_COND_TIMESTEP)
video_w = None
audio_w = None
video_rows_t = None
audio_rows_t = None
if denoise_mask is not None:
targets = mask_row_targets(denoise_mask[0, 0].to(torch.float32), latent_t, lat_h, lat_w)
if targets is not None:
if bool(targets.any()):
video_w = targets.to(torch.float32).unsqueeze(1) # [n, 1], 1 = generate
m = mask_row_values(denoise_mask[0, 0].to(torch.float32), latent_t, lat_h, lat_w)
if m is not None:
rows_t = (1.0 - m * sigma_v.to(m.device)).clamp(max=t_pin_v)
if rows_t.unique().numel() == 1:
seg_t["video"] = float(rows_t[0])
else:
seg_t["video"] = t_pin_v
video_rows_t = rows_t
if audio_denoise_mask is not None:
targets = audio_denoise_mask[0, 0].to(torch.float32).reshape(-1) >= 0.5
if not bool(targets.all()):
if bool(targets.any()):
audio_w = targets.to(torch.float32).unsqueeze(1)
m = audio_denoise_mask[0, 0].to(torch.float32).reshape(-1)
if not bool((m >= 1.0 - 1e-3).all()):
sigma_a = 1.0 - t_a
rows_t = (1.0 - m * sigma_a).clamp(max=t_pin_a)
if rows_t.unique().numel() == 1:
seg_t["audio"] = float(rows_t[0])
else:
seg_t["audio"] = t_pin_a
audio_rows_t = rows_t
unique_t = sorted({t_v, t_a} | {seg_t[k] for _, _, k in layout.segments}
| ({t_pin_v} if video_w is not None else set())
| ({t_pin_a} if audio_w is not None else set()))
| (set(video_rows_t.unique().tolist()) if video_rows_t is not None else set())
| (set(audio_rows_t.unique().tolist()) if audio_rows_t is not None else set()))
t_row = {t: i for i, t in enumerate(unique_t)}
seg_tag = {"text": 1, "video": 0, "audio": 2, "cond": 0, "ref_img": 0, "ref_audio": 2}
def rows_to_mod_index(rows_t, tag):
# per-row timestep values -> per-row mod-row indices into the t_emb table
levels = rows_t.unique()
base = torch.tensor([t_row[v] * 3 + tag for v in levels.tolist()],
dtype=torch.long, device=rows_t.device)
return base[torch.searchsorted(levels, rows_t)]
text_tags = payload.get("text_token_tags")
mod_segments = []
for a, b, kind in layout.segments:
@ -604,10 +613,10 @@ class MiniMaxH3Model(nn.Module):
if i == b - a or tags[i] != tags[run_start]:
mod_segments.append((a + run_start, a + i, row_base + int(tags[run_start])))
run_start = i
elif kind == "video" and video_w is not None:
mod_segments.append((a, b, (row_base + seg_tag[kind], t_row[t_pin_v] * 3 + seg_tag[kind], video_w)))
elif kind == "audio" and audio_w is not None:
mod_segments.append((a, b, (row_base + seg_tag[kind], t_row[t_pin_a] * 3 + seg_tag[kind], audio_w)))
elif kind == "video" and video_rows_t is not None:
mod_segments.append((a, b, rows_to_mod_index(video_rows_t, seg_tag[kind])))
elif kind == "audio" and audio_rows_t is not None:
mod_segments.append((a, b, rows_to_mod_index(audio_rows_t, seg_tag[kind])))
else:
mod_segments.append((a, b, row_base + seg_tag[kind]))
@ -686,12 +695,12 @@ class MiniMaxH3Model(nn.Module):
# target streams are single contiguous segments (audio then video, last two)
va, vb, _ = next(s for s in layout.segments if s[2] == "video")
aa, ab, _ = next(s for s in layout.segments if s[2] == "audio")
if video_w is not None:
video_seg = (va, vb, (t_row[seg_t["video"]], t_row[t_pin_v], video_w))
if video_rows_t is not None:
video_seg = (va, vb, rows_to_mod_index(video_rows_t, 0) // 3)
else:
video_seg = (va, vb, t_row[seg_t["video"]])
if audio_w is not None:
audio_seg = (aa, ab, (t_row[seg_t["audio"]], t_row[t_pin_a], audio_w))
if audio_rows_t is not None:
audio_seg = (aa, ab, rows_to_mod_index(audio_rows_t, 0) // 3)
else:
audio_seg = (aa, ab, t_row[seg_t["audio"]])
v, a = self.final_layer(h, t_emb, video_seg, audio_seg)

30
comfy/model_base.py
View File

@ -2177,9 +2177,9 @@ class MiniMaxH3(BaseModel):
denoise_mask = kwargs.get("denoise_mask", None)
if denoise_mask is not None and latent_shapes is not None and len(latent_shapes) > 1:
masks = utils.unpack_latents(denoise_mask, latent_shapes)
if torch.amin(masks[0]).item() < 0.5:
if torch.amin(masks[0]).item() < 1.0 - 1e-3:
out['denoise_mask'] = comfy.conds.CONDRegular(masks[0][:1, :1].clone())
if torch.amin(masks[1]).item() < 0.5:
if torch.amin(masks[1]).item() < 1.0 - 1e-3:
out['audio_denoise_mask'] = comfy.conds.CONDRegular(masks[1][:1, :1].clone())
if cross_attn is not None and latent_shapes is not None and len(latent_shapes) > 1:
@ -2193,17 +2193,19 @@ class MiniMaxH3(BaseModel):
return out
def process_denoise_mask(self, denoise_masks):
# snap the video mask to the DiT patch grid and the audio mask to whole latent
# frames so a row's timestep label matches its content
vm = denoise_masks[0]
h, w = vm.shape[-2:]
# snap the video mask to the DiT patch grid (2x2 latent pixels per patch)
# and the audio mask to each audio latent frame (which run at 40 audio frames per second)
video_mask = denoise_masks[0]
h, w = video_mask.shape[-2:]
ph, pw = self.diffusion_model.patch_size[1:]
vm = torch.nn.functional.pad(vm, (0, -w % pw, 0, -h % ph))
vm = (vm.reshape(vm.shape[:-2] + (vm.shape[-2] // ph, ph, vm.shape[-1] // pw, pw)).amax(dim=(-3, -1)) >= 0.5).to(vm.dtype)
denoise_masks[0] = vm.repeat_interleave(ph, dim=-2).repeat_interleave(pw, dim=-1)[..., :h, :w]
lead = video_mask.shape[:-2]
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))
denoise_masks[0] = video_mask.repeat_interleave(ph, dim=-2).repeat_interleave(pw, dim=-1)[..., :h, :w]
if len(denoise_masks) > 1:
am = denoise_masks[1]
denoise_masks[1] = (am.amax(dim=1, keepdim=True) >= 0.5).to(am.dtype).expand_as(am).contiguous()
audio_mask = denoise_masks[1].amax(dim=1, keepdim=True)
denoise_masks[1] = audio_mask.expand_as(denoise_masks[1]).contiguous()
return denoise_masks
def scale_latent_inpaint(self, sigma, noise, latent_image, **kwargs):
@ -2213,15 +2215,15 @@ class MiniMaxH3(BaseModel):
return super().scale_latent_inpaint(sigma=sigma, noise=noise, latent_image=latent_image, **kwargs)
cleans = utils.unpack_latents(latent_image, shapes)
noises = utils.unpack_latents(noise, shapes)
aug = comfy.ldm.minimax.model.VISUAL_COND_TIMESTEP
aug = comfy.ldm.minimax.model.VISUAL_COND_TIMESTEP # H3's video timestep is 0.999 by default
cleans[0] = aug * cleans[0] + (1.0 - aug) * noises[0]
scale = self.audio_scale()
if scale != 1.0:
# the sampler carries audio as (sigma_v / sigma_a) * x_audio and latent_image
# holds audio_scale * x_audio, so rescale for the model to see it clean
ms = self.model_sampling
model_sampling = self.model_sampling
sigma_v = sigma.clamp(min=1e-6)
sigma_a = comfy.ldm.minimax.model.time_shift_sigma(sigma_v, ms.shift, ms.audio_shift)
sigma_a = comfy.ldm.minimax.model.time_shift_sigma(sigma_v, model_sampling.shift, model_sampling.audio_shift)
factor = (sigma_v / sigma_a) / scale
cleans[1] = cleans[1] * factor.view(factor.shape[:1] + (1,) * (cleans[1].ndim - 1)).to(cleans[1].dtype)
return utils.pack_latents(cleans)[0]