From 68bcf109ffbf0b020fdc880b7af2bef8404ce134 Mon Sep 17 00:00:00 2001 From: digital-garbage <133010426+digital-garbage@users.noreply.github.com> Date: Tue, 25 Aug 2026 22:33:18 +0300 Subject: [PATCH] fix(h3): the token-refiner edit renamed the checkpoint's own weight keys MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Reported from a rental: manual gain mode with the refinement keys deleted died with "'TokenRefiner' object has no attribute 'inner'". Learned mode was fine, because with no key the gains are neutral and the installer returns before building anything — so only manual, which ignores keys, ever built the wrapper. Wrapping a module in an nn.Module renames every weight beneath it (token_refiner.blocks.* -> token_refiner.inner.blocks.*). patch_model applies object patches BEFORE loading weights, so the names were recorded WITH .inner; unpatch_model restores weights BEFORE unwrapping, and with several patchers alive off one shared model — one clone per scene, plus one for the AdaLN gains — another patcher unwraps first and the restore walks .inner on a real TokenRefiner. Both wrappers now patch `.forward` through the same object-patch mechanism instead of replacing the module, so the tree and the state dict stay exactly as the checkpoint built them. Neither is an nn.Module any more, and each captures the original bound method rather than re-reading `.forward` at call time — ComfyUI restores a patch by writing the previous value back, which can be a live wrapper, so re-reading would nest a scale per scene. Reproduced against faithful patch semantics before and after. AdaLN gains had the identical latent hazard and are fixed the same way. --- minimax_h3.py | 83 +++++++++++++++++++++++----- tests/test_h3_token_refiner_edit.py | 85 ++++++++++++++++++++++++++++- 2 files changed, 152 insertions(+), 16 deletions(-) diff --git a/minimax_h3.py b/minimax_h3.py index a548ca52..c187fa53 100644 --- a/minimax_h3.py +++ b/minimax_h3.py @@ -1328,16 +1328,56 @@ _ADALN_GATE_CHUNKS = (2, 5) # shift_msa scale_msa GATE_msa shift_mlp sc _ADALN_MODALITIES = 3 -class AdalnModalityGain(torch.nn.Module): - """Wraps one block's ``adaln_proj`` and scales its gates per modality.""" +def unwrapped_forward(module): + """The module's own ``forward``, never a FunPack wrapper already sitting on it. + + ComfyUI restores an object patch by writing the previous value back, and the previous + value of ``.forward`` may itself be one of our wrappers from a run that is still + winding down. Capturing that would nest wrappers on every scene and eventually stack a + scale per scene, so the chain is followed back to the real bound method every time. + """ + fn = getattr(module, "forward", module) + seen = 0 + while getattr(fn, "_funpack_wrapper", False) and seen < 16: + fn = getattr(fn, "original", fn) + seen += 1 + return fn + + +def patch_module_forward(patcher, key, wrapper): + """Install `wrapper` as `.forward` through ComfyUI's own object-patch mechanism. + + Patching the FORWARD rather than replacing the module is what keeps the state dict + intact: a wrapper module would rename every weight under `key` to `.inner.*`, and + ComfyUI captures weight names off the live tree, so an unwrap by any other patcher + sharing this model would leave a restore walking a `.inner` that no longer exists. + """ + patcher.add_object_patch(f"{key}.forward", wrapper) + + +class AdalnModalityGain: + """Wraps one block's ``adaln_proj.forward`` and scales its gates per modality. + + Deliberately NOT an ``nn.Module``. Replacing a module with a wrapper module renames + every weight beneath it in the state dict (``adaln_proj.weight`` becomes + ``adaln_proj.inner.weight``), and ComfyUI records weight keys off the LIVE tree while + the patch is applied — so an unwrap by any other patcher sharing this model leaves a + restore walking a ``.inner`` that is no longer there. Patching ``forward`` leaves the + tree and the state dict exactly as the checkpoint built them. + """ + + #: marks an already-installed wrapper, so a re-install unwraps instead of nesting + _funpack_wrapper = True def __init__(self, inner, gains): - super().__init__() self.inner = inner + # The ORIGINAL callable, captured before the patch is applied. Calling + # ``inner.forward`` at call time would re-enter this wrapper once it is installed. + self.original = unwrapped_forward(inner) self.gains = {int(t): float(g) for t, g in (gains or {}).items() if float(g) != 1.0} - def forward(self, t_emb): - out = self.inner(t_emb) + def __call__(self, t_emb): + out = self.original(t_emb) if not self.gains or getattr(self.inner, "modalities", 1) != _ADALN_MODALITIES: return out out = list(out) @@ -1378,7 +1418,7 @@ def apply_adaln_gains(model, gains): inner = patched.get_model_object(key) except Exception: continue - patched.add_object_patch(key, AdalnModalityGain(inner, live)) + patch_module_forward(patched, key, AdalnModalityGain(inner, live)) count += 1 if not count: return model, "AdaLN gain skipped — no adaln_proj found on the blocks." @@ -1411,13 +1451,28 @@ def apply_adaln_gains(model, gains): # head of this span and are not ours to steer. -class TokenRefinerEdit(torch.nn.Module): - """Wraps ``token_refiner`` and edits its output over a row range.""" +class TokenRefinerEdit: + """Wraps ``token_refiner.forward`` and edits its output over a row range. + + Deliberately NOT an ``nn.Module``. Replacing the module with a wrapper module renames + the checkpoint's own keys — ``token_refiner.blocks.*`` becomes + ``token_refiner.inner.blocks.*`` — and ComfyUI's ``patch_model`` applies object patches + BEFORE it loads weights, so every weight gets recorded under the ``.inner`` name. On the + way out it restores weights before it unwraps, and with several patchers alive off one + shared model (one clone per scene, plus one for the AdaLN gains) another patcher can + unwrap first — leaving a restore that walks ``.inner`` on a real ``TokenRefiner``. + Patching the forward leaves the tree and the state dict as the checkpoint built them. + """ + + #: marks an already-installed wrapper, so a re-install unwraps instead of nesting + _funpack_wrapper = True def __init__(self, inner, scale=1.0, bias=None, row_start=0, row_end=None, bias_relative=False): - super().__init__() self.inner = inner + # The ORIGINAL callable, captured before the patch is applied. Calling + # ``inner.forward`` at call time would re-enter this wrapper once installed. + self.original = unwrapped_forward(inner) self.scale = float(scale) self.bias = bias self.row_start = int(row_start) @@ -1430,8 +1485,8 @@ class TokenRefinerEdit(torch.nn.Module): # prompt row", which is the same thing on every prompt and every checkpoint. self.bias_relative = bool(bias_relative) - def forward(self, x, *args, **kwargs): - out = self.inner(x, *args, **kwargs) + def __call__(self, x, *args, **kwargs): + out = self.original(x, *args, **kwargs) try: rows = out.shape[-2] start = max(0, min(int(self.row_start), rows)) @@ -1450,7 +1505,7 @@ class TokenRefinerEdit(torch.nn.Module): span = span + bias out[..., start:end, :] = span except Exception: # a refiner shape we do not recognise: leave it alone - return self.inner(x, *args, **kwargs) + return self.original(x, *args, **kwargs) return out @@ -1490,8 +1545,8 @@ def apply_token_refiner_edit(model, scale=1.0, bias=None, row_start=0, row_end=N except Exception as error: return model, f"token-refiner edit skipped — no token_refiner to patch ({error})." patched = model.clone() - patched.add_object_patch(key, TokenRefinerEdit(inner, scale, bias, row_start, row_end, - bias_relative=bias_relative)) + patch_module_forward(patched, key, TokenRefinerEdit(inner, scale, bias, row_start, + row_end, bias_relative=bias_relative)) parts = [] if float(scale) != 1.0: parts.append(f"scale={float(scale):g}") diff --git a/tests/test_h3_token_refiner_edit.py b/tests/test_h3_token_refiner_edit.py index 7a1d93e3..71be1f74 100644 --- a/tests/test_h3_token_refiner_edit.py +++ b/tests/test_h3_token_refiner_edit.py @@ -121,13 +121,36 @@ def test_no_edit_returns_the_model_unpatched(): assert out is model and note is None and model.patched == {} -def test_it_patches_the_refiner_object(): +def test_it_patches_the_forward_not_the_module(): + """The module itself must stay in place: swapping it for a wrapper MODULE renames every + weight under it (token_refiner.blocks.* -> token_refiner.inner.blocks.*), and ComfyUI + records weight names off the live tree while the patch is applied. An unwrap by any + other patcher sharing the model then leaves a restore walking a `.inner` that is gone — + 'TokenRefiner object has no attribute inner', mid-sample.""" model = FakePatcher() out, note = h3.apply_token_refiner_edit(model, scale=1.2, row_start=4) - assert isinstance(out.patched["diffusion_model.token_refiner"], h3.TokenRefinerEdit) + assert "diffusion_model.token_refiner" not in out.patched + assert isinstance(out.patched["diffusion_model.token_refiner.forward"], h3.TokenRefinerEdit) assert "scale=1.2" in note and "rows 4:end" in note +def test_the_edit_is_not_a_module_so_the_state_dict_is_untouched(): + assert not isinstance(h3.TokenRefinerEdit(FakeRefiner()), torch.nn.Module) + assert not isinstance(h3.AdalnModalityGain(FakeRefiner(), {}), torch.nn.Module) + + +def test_reinstalling_unwraps_instead_of_nesting(): + """Per scene we install again off a shared model. ComfyUI restores an object patch by + writing the previous value back, so the previous `forward` can be a live wrapper — and + capturing that would stack another scale on every scene.""" + inner = FakeRefiner() + first = h3.TokenRefinerEdit(inner, scale=2.0) + inner.forward = first # as patch_model would leave it + second = h3.TokenRefinerEdit(inner, scale=2.0) + assert second.original is not first + assert torch.allclose(second(torch.zeros(1, 4, 4)), torch.full((4, HID), 2.0)) + + def test_object_patches_so_nothing_outlives_the_run(): src = inspect.getsource(h3.apply_token_refiner_edit) assert "add_object_patch" in src and "model.clone()" in src @@ -248,3 +271,61 @@ def test_the_note_says_the_bias_is_relative(): model = FakePatcher() _out, note = h3.apply_token_refiner_edit(model, bias=torch.zeros(HID), bias_relative=True) assert "bias (relative)" in note + + +# ── the state-dict hazard that crashed a real run ─────────────────────────── +# +# Reported from a rental: manual gain mode with the refinement keys deleted, sampling dies +# with "'TokenRefiner' object has no attribute 'inner'". Learned mode was fine, because with +# no key the gains are neutral and the installer returns before building anything. +# +# The chain: a wrapper MODULE renames every weight under it, ComfyUI's patch_model applies +# object patches BEFORE loading weights (so names are recorded WITH `.inner`), and +# unpatch_model restores weights BEFORE unwrapping. With more than one patcher alive off one +# shared model — one clone per scene, plus one for the AdaLN gains — another patcher can +# unwrap first, and the weight restore then walks `.inner` on a real TokenRefiner. + +class _Refiner(torch.nn.Module): + def __init__(self): + super().__init__() + self.blocks = torch.nn.Linear(4, 4) + + def forward(self, x, *a, **k): + return torch.ones(x.shape[-2], HID) + + +class _DiT(torch.nn.Module): + def __init__(self): + super().__init__() + self.token_refiner = _Refiner() + + +class _Base(torch.nn.Module): + def __init__(self): + super().__init__() + self.diffusion_model = _DiT() + + +def test_the_patch_never_renames_a_checkpoint_weight(): + """The whole bug in one assertion: no weight key may gain an `.inner` segment.""" + model = _Base() + before = list(dict(model.named_parameters())) + edit = h3.TokenRefinerEdit(model.diffusion_model.token_refiner, scale=1.2) + model.diffusion_model.token_refiner.forward = edit # as patch_model installs it + assert list(dict(model.named_parameters())) == before + assert not any(".inner." in k for k in dict(model.named_parameters())) + + +def test_the_module_itself_stays_in_place(): + model = _Base() + original = model.diffusion_model.token_refiner + edit = h3.TokenRefinerEdit(original, scale=1.2) + model.diffusion_model.token_refiner.forward = edit + assert model.diffusion_model.token_refiner is original + assert isinstance(model.diffusion_model.token_refiner, _Refiner) + + +def test_the_adaln_gain_patches_a_forward_too(): + """Same hazard, same shape: adaln_proj.weight must not become adaln_proj.inner.weight.""" + src = inspect.getsource(h3.apply_adaln_gains) + assert "patch_module_forward" in src