zlei6/ComfyUI
1
0
Fork 0
mirror of https://github.com/Comfy-Org/ComfyUI.git synced 2026-09-07 18:00:11 +08:00
Code Issues Packages Projects Releases Wiki Activity

MiniMax-H3: Support PDD LoRA (#15908)

Browse Source
This commit is contained in:
Jukka Seppänen
2026-08-29 01:12:36 +03:00
committed by GitHub
parent 3ac5d7941d
commit 2504e68d4d
2 changed files with 33 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -303,7 +303,7 @@ class FinalLayer(nn.Module):
self.video_out = operations.Linear(hidden, video_dim, bias=True, dtype=torch.float32, device=device)
self.audio_out = operations.Linear(hidden, audio_dim, bias=True, dtype=torch.float32, device=device)
def forward(self, x, t_emb, video_seg, audio_seg):
def forward(self, x, t_emb, video_seg, audio_seg, sigma, sample_sigmas, shifts):
# video_seg / audio_seg: (start, stop, row) of the target streams, where row
# is a mod-row index or a per-token blend (see _mod_row)
shift, scale = self.adaln_proj(t_emb)
@@ -312,7 +312,33 @@ class FinalLayer(nn.Module):
a, b, row = seg
return (self.norm(x[a:b]) * (1.0 + _mod_row(scale, row, scale.dtype)) + _mod_row(shift, row, shift.dtype)).to(torch.float32)
return self.video_out(mod(video_seg)), self.audio_out(mod(audio_seg))
n = self.video_out.weight.shape[0] // self.video_out.out_features
if n == 1:
return self.video_out(mod(video_seg)), self.audio_out(mod(audio_seg))
# PDD head bank: row block 0 is a full head, later blocks are offsets from it;
# a step consumes the dt-weighted mean of the heads it spans.
if sample_sigmas is None:
raise ValueError("MiniMax H3 PDD heads need the sampler's sigma schedule")
i = int((sample_sigmas - sigma).abs().argmin())
sigma_next = sample_sigmas[min(i + 1, sample_sigmas.shape[0] - 1)]
start, stop = (round(float(1.0 - time_shift_sigma(s, shifts[0], 1.0)) * n) for s in (sigma, sigma_next))
start = min(start, n - 1)
stop = max(stop, start + 1)
return (_pdd_head(self.video_out, mod(video_seg), n, start, stop, shifts[0]),
_pdd_head(self.audio_out, mod(audio_seg), n, start, stop, shifts[1]))
def _pdd_head(head, h, n, start, stop, flow_shift):
grid = torch.linspace(1.0, 0.0, n + 1, dtype=torch.float64)
dt = (1.0 - flow_shift * grid / (1.0 + (flow_shift - 1.0) * grid)).diff()[start:stop]
w = (dt / dt.sum()).to(h)
with comfy.ops.CastBiasWeightContext(head, h, offloadable=True) as (weight, bias):
rows = weight.reshape(n, -1, weight.shape[1])
brows = bias.reshape(n, -1)
first = max(start, 1)
return nn.functional.linear(h, rows[0] + torch.einsum("n,noi->oi", w[first - start:], rows[first:stop]),
brows[0] + torch.einsum("n,no->o", w[first - start:], brows[first:stop]))
class PackedLayout:
@@ -723,7 +749,7 @@ class MiniMaxH3Model(nn.Module):
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)
v, a = self.final_layer(h, t_emb, video_seg, audio_seg, sigma_v, transformer_options.get("sample_sigmas"), (shift_v, shift_a))
video_out = unpatchify_video(v, latent_t, lat_h // 2, lat_w // 2, self.latents_dim, self.patch_size)
video_out = video_out[:, :, :orig_t, :orig_h, :orig_w]

8
comfy/model_patcher.py
View File

@@ -1020,14 +1020,14 @@ class ModelPatcher:
m.bias_function = []
if weight_key in self.patches:
if force_patch_weights:
if force_patch_weights or comfy.lora.calculate_shape(self.patches[weight_key], m.weight, weight_key) != m.weight.shape:
self.patch_weight_to_device(weight_key)
else:
_, set_func, convert_func = get_key_weight(self.model, weight_key)
m.weight_function = [LowVramPatch(weight_key, self.patches, convert_func, set_func)]
patch_counter += 1
if bias_key in self.patches:
if force_patch_weights:
if force_patch_weights or comfy.lora.calculate_shape(self.patches[bias_key], m.bias, bias_key) != m.bias.shape:
self.patch_weight_to_device(bias_key)
else:
_, set_func, convert_func = get_key_weight(self.model, bias_key)
@@ -1217,14 +1217,14 @@ class ModelPatcher:
module_mem += move_weight_functions(m, device_to)
if lowvram_possible:
if weight_key in self.patches:
if force_patch_weights:
if force_patch_weights or comfy.lora.calculate_shape(self.patches[weight_key], m.weight, weight_key) != m.weight.shape:
self.patch_weight_to_device(weight_key)
else:
_, set_func, convert_func = get_key_weight(self.model, weight_key)
m.weight_function.append(LowVramPatch(weight_key, self.patches, convert_func, set_func))
patch_counter += 1
if bias_key in self.patches:
if force_patch_weights:
if force_patch_weights or comfy.lora.calculate_shape(self.patches[bias_key], m.bias, bias_key) != m.bias.shape:
self.patch_weight_to_device(bias_key)
else:
_, set_func, convert_func = get_key_weight(self.model, bias_key)