feat: PTD few-step support for MiniMax H3 as a module
Settings > Engine > Sampling: pick a PTD adapter file from models/ptd and the run samples the file's own 4 to 8 hops with no CFG, whatever sampler and steps are chosen. Reads the EPFL fl2va and taobao-content-ai Ref2VA releases (bf16 or int8, full or pruned base); ordinary LoRAs stack. Ported from ComfyUI-PTD-Compact (Nynxz), ComfyUI-PTD (EPFL VITA Lab) and ComfyUI-MiniMaxH3-PTD (taobao-content-ai), all Apache-2.0: README Thank you plus NOTICE and licence text in the module. Proven on a tiny real H3 on CPU only. Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_0164nckPWA2L2ugVqiDcXzc3
This commit is contained in:
Claude
@@ -1,5 +1,11 @@
|
||||
# Changelog
|
||||
|
||||
## [Unreleased]
|
||||
|
||||
### Added (experimental)
|
||||
|
||||
- **PTD few-step** (Settings > Engine > Sampling, MiniMax H3 only): renders in the 4 to 8 hops of a PTD adapter file from `ComfyUI/models/ptd` instead of a long schedule, with no CFG. Reads the EPFL fl2va and taobao-content-ai Ref2VA releases, as bf16 or int8, on the full or pruned base; ordinary LoRAs stack with it. The chosen sampler, steps and scheduler are replaced while it is on. Ported from ComfyUI-PTD-Compact (Nynxz). Not yet run on a real H3.
|
||||
|
||||
## [5.0.1] - 2026-10-21
|
||||
|
||||
Fixes from real H3 runs on rentals, plus the first usable LoRA training.
|
||||
|
||||
@@ -160,6 +160,7 @@ I express my deepest gratitude to:
|
||||
- [OpenCut](https://github.com/opencut-app/opencut-classic) — the in-browser non-linear video editor whose UI and interaction patterns inspired FunPack's timeline studio;
|
||||
- [Inter](https://github.com/rsms/inter) and IBM Plex Mono — the fonts of the app (SIL Open Font License; copies in `app/assets/fonts`);
|
||||
- [musubi-tuner](https://github.com/kohya-ss/musubi-tuner) by kohya-ss (Apache-2.0) — the LoRA training add-on's loop (which layers get a LoRA, per-block checkpointing, noise sampling and loss weighting, optimizer and learning rate) and its block swap (a number of transformer blocks kept in system memory) follow its MiniMax H3 training recipe;
|
||||
- PTD few-step support for MiniMax H3 is ported from [ComfyUI-PTD-Compact](https://github.com/Nynxz/ComfyUI-PTD-Compact) by Nynxz (Apache-2.0), which builds on [ComfyUI-PTD](https://github.com/vita-epfl/PTD-inference) by EPFL VITA Lab (Apache-2.0; paper "Parametric Trajectory Distillation for Few-Step Video Generation", arXiv 2610.11498) and on ComfyUI-MiniMaxH3-PTD by taobao-content-ai (Apache-2.0); licence notice in `modules/sampling/ptd`;
|
||||
- ComfyUI and its whole community.
|
||||
|
||||
Without all of you, this project would've been impossible.
|
||||
|
||||
@@ -67,8 +67,9 @@ def _finite(out):
|
||||
|
||||
|
||||
#: Failures that come from the stand-in model, not from the module: a four-block model has no blocks 31-40, and the
|
||||
#: LTX stand-in is not a model_base class with the windowing hook (the module's own tests patch that in).
|
||||
TINY_MODEL = ("no usable blocks", "map_context_window_to_modalities")
|
||||
#: LTX stand-in is not a model_base class with the windowing hook (the module's own tests patch that in). Likewise
|
||||
#: PTD has no adapter file on a test machine (its own tests build one).
|
||||
TINY_MODEL = ("no usable blocks", "map_context_window_to_modalities", "no PTD adapter file")
|
||||
|
||||
|
||||
def _install_failures(status):
|
||||
|
||||
@@ -0,0 +1,202 @@
|
||||
|
||||
Apache License
|
||||
Version 2.0, January 2004
|
||||
http://www.apache.org/licenses/
|
||||
|
||||
TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
|
||||
|
||||
1. Definitions.
|
||||
|
||||
"License" shall mean the terms and conditions for use, reproduction,
|
||||
and distribution as defined by Sections 1 through 9 of this document.
|
||||
|
||||
"Licensor" shall mean the copyright owner or entity authorized by
|
||||
the copyright owner that is granting the License.
|
||||
|
||||
"Legal Entity" shall mean the union of the acting entity and all
|
||||
other entities that control, are controlled by, or are under common
|
||||
control with that entity. For the purposes of this definition,
|
||||
"control" means (i) the power, direct or indirect, to cause the
|
||||
direction or management of such entity, whether by contract or
|
||||
otherwise, or (ii) ownership of fifty percent (50%) or more of the
|
||||
outstanding shares, or (iii) beneficial ownership of such entity.
|
||||
|
||||
"You" (or "Your") shall mean an individual or Legal Entity
|
||||
exercising permissions granted by this License.
|
||||
|
||||
"Source" form shall mean the preferred form for making modifications,
|
||||
including but not limited to software source code, documentation
|
||||
source, and configuration files.
|
||||
|
||||
"Object" form shall mean any form resulting from mechanical
|
||||
transformation or translation of a Source form, including but
|
||||
not limited to compiled object code, generated documentation,
|
||||
and conversions to other media types.
|
||||
|
||||
"Work" shall mean the work of authorship, whether in Source or
|
||||
Object form, made available under the License, as indicated by a
|
||||
copyright notice that is included in or attached to the work
|
||||
(an example is provided in the Appendix below).
|
||||
|
||||
"Derivative Works" shall mean any work, whether in Source or Object
|
||||
form, that is based on (or derived from) the Work and for which the
|
||||
editorial revisions, annotations, elaborations, or other modifications
|
||||
represent, as a whole, an original work of authorship. For the purposes
|
||||
of this License, Derivative Works shall not include works that remain
|
||||
separable from, or merely link (or bind by name) to the interfaces of,
|
||||
the Work and Derivative Works thereof.
|
||||
|
||||
"Contribution" shall mean any work of authorship, including
|
||||
the original version of the Work and any modifications or additions
|
||||
to that Work or Derivative Works thereof, that is intentionally
|
||||
submitted to Licensor for inclusion in the Work by the copyright owner
|
||||
or by an individual or Legal Entity authorized to submit on behalf of
|
||||
the copyright owner. For the purposes of this definition, "submitted"
|
||||
means any form of electronic, verbal, or written communication sent
|
||||
to the Licensor or its representatives, including but not limited to
|
||||
communication on electronic mailing lists, source code control systems,
|
||||
and issue tracking systems that are managed by, or on behalf of, the
|
||||
Licensor for the purpose of discussing and improving the Work, but
|
||||
excluding communication that is conspicuously marked or otherwise
|
||||
designated in writing by the copyright owner as "Not a Contribution."
|
||||
|
||||
"Contributor" shall mean Licensor and any individual or Legal Entity
|
||||
on behalf of whom a Contribution has been received by Licensor and
|
||||
subsequently incorporated within the Work.
|
||||
|
||||
2. Grant of Copyright License. Subject to the terms and conditions of
|
||||
this License, each Contributor hereby grants to You a perpetual,
|
||||
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
|
||||
copyright license to reproduce, prepare Derivative Works of,
|
||||
publicly display, publicly perform, sublicense, and distribute the
|
||||
Work and such Derivative Works in Source or Object form.
|
||||
|
||||
3. Grant of Patent License. Subject to the terms and conditions of
|
||||
this License, each Contributor hereby grants to You a perpetual,
|
||||
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
|
||||
(except as stated in this section) patent license to make, have made,
|
||||
use, offer to sell, sell, import, and otherwise transfer the Work,
|
||||
where such license applies only to those patent claims licensable
|
||||
by such Contributor that are necessarily infringed by their
|
||||
Contribution(s) alone or by combination of their Contribution(s)
|
||||
with the Work to which such Contribution(s) was submitted. If You
|
||||
institute patent litigation against any entity (including a
|
||||
cross-claim or counterclaim in a lawsuit) alleging that the Work
|
||||
or a Contribution incorporated within the Work constitutes direct
|
||||
or contributory patent infringement, then any patent licenses
|
||||
granted to You under this License for that Work shall terminate
|
||||
as of the date such litigation is filed.
|
||||
|
||||
4. Redistribution. You may reproduce and distribute copies of the
|
||||
Work or Derivative Works thereof in any medium, with or without
|
||||
modifications, and in Source or Object form, provided that You
|
||||
meet the following conditions:
|
||||
|
||||
(a) You must give any other recipients of the Work or
|
||||
Derivative Works a copy of this License; and
|
||||
|
||||
(b) You must cause any modified files to carry prominent notices
|
||||
stating that You changed the files; and
|
||||
|
||||
(c) You must retain, in the Source form of any Derivative Works
|
||||
that You distribute, all copyright, patent, trademark, and
|
||||
attribution notices from the Source form of the Work,
|
||||
excluding those notices that do not pertain to any part of
|
||||
the Derivative Works; and
|
||||
|
||||
(d) If the Work includes a "NOTICE" text file as part of its
|
||||
distribution, then any Derivative Works that You distribute must
|
||||
include a readable copy of the attribution notices contained
|
||||
within such NOTICE file, excluding those notices that do not
|
||||
pertain to any part of the Derivative Works, in at least one
|
||||
of the following places: within a NOTICE text file distributed
|
||||
as part of the Derivative Works; within the Source form or
|
||||
documentation, if provided along with the Derivative Works; or,
|
||||
within a display generated by the Derivative Works, if and
|
||||
wherever such third-party notices normally appear. The contents
|
||||
of the NOTICE file are for informational purposes only and
|
||||
do not modify the License. You may add Your own attribution
|
||||
notices within Derivative Works that You distribute, alongside
|
||||
or as an addendum to the NOTICE text from the Work, provided
|
||||
that such additional attribution notices cannot be construed
|
||||
as modifying the License.
|
||||
|
||||
You may add Your own copyright statement to Your modifications and
|
||||
may provide additional or different license terms and conditions
|
||||
for use, reproduction, or distribution of Your modifications, or
|
||||
for any such Derivative Works as a whole, provided Your use,
|
||||
reproduction, and distribution of the Work otherwise complies with
|
||||
the conditions stated in this License.
|
||||
|
||||
5. Submission of Contributions. Unless You explicitly state otherwise,
|
||||
any Contribution intentionally submitted for inclusion in the Work
|
||||
by You to the Licensor shall be under the terms and conditions of
|
||||
this License, without any additional terms or conditions.
|
||||
Notwithstanding the above, nothing herein shall supersede or modify
|
||||
the terms of any separate license agreement you may have executed
|
||||
with Licensor regarding such Contributions.
|
||||
|
||||
6. Trademarks. This License does not grant permission to use the trade
|
||||
names, trademarks, service marks, or product names of the Licensor,
|
||||
except as required for reasonable and customary use in describing the
|
||||
origin of the Work and reproducing the content of the NOTICE file.
|
||||
|
||||
7. Disclaimer of Warranty. Unless required by applicable law or
|
||||
agreed to in writing, Licensor provides the Work (and each
|
||||
Contributor provides its Contributions) on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
|
||||
implied, including, without limitation, any warranties or conditions
|
||||
of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
|
||||
PARTICULAR PURPOSE. You are solely responsible for determining the
|
||||
appropriateness of using or redistributing the Work and assume any
|
||||
risks associated with Your exercise of permissions under this License.
|
||||
|
||||
8. Limitation of Liability. In no event and under no legal theory,
|
||||
whether in tort (including negligence), contract, or otherwise,
|
||||
unless required by applicable law (such as deliberate and grossly
|
||||
negligent acts) or agreed to in writing, shall any Contributor be
|
||||
liable to You for damages, including any direct, indirect, special,
|
||||
incidental, or consequential damages of any character arising as a
|
||||
result of this License or out of the use or inability to use the
|
||||
Work (including but not limited to damages for loss of goodwill,
|
||||
work stoppage, computer failure or malfunction, or any and all
|
||||
other commercial damages or losses), even if such Contributor
|
||||
has been advised of the possibility of such damages.
|
||||
|
||||
9. Accepting Warranty or Additional Liability. While redistributing
|
||||
the Work or Derivative Works thereof, You may choose to offer,
|
||||
and charge a fee for, acceptance of support, warranty, indemnity,
|
||||
or other liability obligations and/or rights consistent with this
|
||||
License. However, in accepting such obligations, You may act only
|
||||
on Your own behalf and on Your sole responsibility, not on behalf
|
||||
of any other Contributor, and only if You agree to indemnify,
|
||||
defend, and hold each Contributor harmless for any liability
|
||||
incurred by, or claims asserted against, such Contributor by reason
|
||||
of your accepting any such warranty or additional liability.
|
||||
|
||||
END OF TERMS AND CONDITIONS
|
||||
|
||||
APPENDIX: How to apply the Apache License to your work.
|
||||
|
||||
To apply the Apache License to your work, attach the following
|
||||
boilerplate notice, with the fields enclosed by brackets "[]"
|
||||
replaced with your own identifying information. (Don't include
|
||||
the brackets!) The text should be enclosed in the appropriate
|
||||
comment syntax for the file format. We also recommend that a
|
||||
file or class name and description of purpose be included on the
|
||||
same "printed page" as the copyright notice for easier
|
||||
identification within third-party archives.
|
||||
|
||||
Copyright [yyyy] [name of copyright owner]
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
You may obtain a copy of the License at
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
@@ -0,0 +1,14 @@
|
||||
PTD few-step support for MiniMax-H3 (modules/sampling/ptd)
|
||||
|
||||
This folder is a port of existing Apache-2.0 work; the licence text is in LICENSE-Apache-2.0.txt.
|
||||
|
||||
* ComfyUI-PTD-Compact by Nynxz (https://github.com/Nynxz/ComfyUI-PTD-Compact): the compact int8 / reduced-rank
|
||||
adapter controller, the in-place residual, the DynamicVRAM handling, the Ref2VA support and the tests.
|
||||
* ComfyUI-PTD (https://github.com/vita-epfl/PTD-inference), Copyright (c) 2026 EPFL VITA Lab: the hop
|
||||
controller, forward hooks, block arithmetic, audio clock, checkpoint format and fixed-hop sampler, from the
|
||||
paper "Parametric Trajectory Distillation for Few-Step Video Generation" (arXiv:2610.11498).
|
||||
* ComfyUI-MiniMaxH3-PTD by taobao-content-ai
|
||||
(https://huggingface.co/taobao-content-ai/MiniMax-H3-PTD-LoRAs): the Ref2VA hop grid, output-head bank
|
||||
fusion, audio chord correction and stacked-hop key mapping.
|
||||
|
||||
The adapter weights and MiniMax-H3 have their own licences; see their model repositories.
|
||||
@@ -0,0 +1,137 @@
|
||||
"""PTD few-step adapters for MiniMax-H3: 4 to 8 hops instead of a long schedule, with no CFG.
|
||||
|
||||
A PTD file is not a LoRA: it holds one set of adapters per hop, and the sampler must visit exactly the
|
||||
times it was trained at. Switching this on loads the file, puts the adapters on the model and runs PTD's
|
||||
own sampler over the file's hop times, whatever sampler and steps are chosen. Ordinary LoRAs stack with it.
|
||||
|
||||
Loads the EPFL fl2va releases and the taobao-content-ai Ref2VA releases, as bf16 or as the smaller int8
|
||||
files, on the full or the pruned base. Ported from ComfyUI-PTD-Compact (Nynxz) over ComfyUI-PTD (EPFL VITA
|
||||
Lab) and ComfyUI-MiniMaxH3-PTD (taobao-content-ai), all Apache-2.0; see the README Thank you section.
|
||||
Not yet run on a real H3 here.
|
||||
"""
|
||||
|
||||
import os
|
||||
from pathlib import Path
|
||||
|
||||
import folder_paths
|
||||
from comfy.patcher_extension import WrappersMP
|
||||
|
||||
from ..._core import control, log, patching
|
||||
|
||||
from . import adapters
|
||||
from .controller import CONTROLLER_KEY, Controller, attach
|
||||
from .sampler import sample_ptd
|
||||
|
||||
ID = "ptd"
|
||||
TITLE = "PTD few-step"
|
||||
MOUNT = "generation.sampling"
|
||||
STAGE = "load"
|
||||
CATEGORY = "sampling"
|
||||
STATUS = "experimental"
|
||||
REQUIRES = ["audio_stream", "adaln_modalities"]
|
||||
|
||||
if "ptd" not in folder_paths.folder_names_and_paths:
|
||||
folder_paths.add_model_folder_path("ptd", str(Path(folder_paths.models_dir) / "ptd"))
|
||||
|
||||
KERNELS = {"exact": False, "fast": True}
|
||||
|
||||
|
||||
def _files():
|
||||
try:
|
||||
return [f for f in folder_paths.get_filename_list("ptd") if f.endswith(".safetensors")]
|
||||
except Exception: # noqa: BLE001 -- no folder yet is the same as no files
|
||||
return []
|
||||
|
||||
|
||||
_FOUND = _files()
|
||||
|
||||
SETTINGS = {
|
||||
"enabled": {
|
||||
"type": "bool", "default": False,
|
||||
"label": "Few-step PTD",
|
||||
"hint": "Renders in the adapter file's own 4 to 8 steps with no CFG. The sampler, steps and scheduler "
|
||||
"are replaced; PTD cannot start from a picture or a mask.",
|
||||
},
|
||||
"adapter": {
|
||||
"type": "enum", "default": _FOUND[0] if _FOUND else "",
|
||||
"options": [{"value": f, "label": f} for f in _FOUND] or [{"value": "", "label": "No file in models/ptd"}],
|
||||
"label": "PTD adapter file",
|
||||
"hint": "A PTD file from ComfyUI/models/ptd. Restart ComfyUI after adding one.",
|
||||
"when": {"enabled": True},
|
||||
},
|
||||
"kernels": {
|
||||
"type": "enum", "default": "exact",
|
||||
"options": [{"value": "exact", "label": "Exact"}, {"value": "fast", "label": "Fast"}],
|
||||
"label": "Block maths",
|
||||
"hint": "Exact follows the arithmetic the adapters were trained with. Fast uses ComfyUI's own blocks, "
|
||||
"about 4 to 8 percent quicker per step; only the rounding differs.",
|
||||
"when": {"enabled": True},
|
||||
},
|
||||
}
|
||||
|
||||
# One parsed file is kept between runs, so changing a seed or a prompt does not read gigabytes again.
|
||||
# It holds no model, so nothing here keeps a diffusion model alive.
|
||||
_kept = {}
|
||||
|
||||
|
||||
def _adapter(path):
|
||||
stamp = (str(path), os.path.getmtime(path), os.path.getsize(path))
|
||||
if _kept.get("stamp") != stamp:
|
||||
_kept.clear()
|
||||
_kept.update(stamp=stamp, adapter=adapters.parse(path, label=Path(path).name))
|
||||
return _kept["adapter"]
|
||||
|
||||
|
||||
def _swap(controller):
|
||||
"""Run PTD's sampler over the adapter's own hop times, in place of the sampler that was chosen."""
|
||||
import comfy.samplers
|
||||
ptd = comfy.samplers.KSAMPLER(sample_ptd, extra_options={"controller": controller})
|
||||
|
||||
def run(executor, guider, sigmas, extra_args, callback, noise, latent_image, denoise_mask, disable_pbar):
|
||||
if abs(float(getattr(guider, "cfg", 1.0)) - 1.0) > 1e-6:
|
||||
log.once("ptd:cfg", log.ALERT, "FunPack PTD",
|
||||
f"CFG is {guider.cfg:g}: PTD was trained without guidance, so set CFG to 1")
|
||||
try:
|
||||
import latent_preview
|
||||
callback = latent_preview.prepare_callback(guider.model_patcher, controller.hops)
|
||||
except Exception: # noqa: BLE001 -- only the progress bar and previews are lost
|
||||
pass
|
||||
log.once("ptd:swap", log.INFO, "FunPack PTD",
|
||||
f"sampling {controller.hops} hops from {controller.adapter.label}; the chosen sampler and steps "
|
||||
f"({len(sigmas) - 1}) are not used")
|
||||
return ptd.sample(guider, controller.grid.float(), extra_args, callback, noise, latent_image,
|
||||
denoise_mask, disable_pbar)
|
||||
|
||||
return run
|
||||
|
||||
|
||||
def install(patcher, values, key):
|
||||
if not values["enabled"]:
|
||||
_kept.clear()
|
||||
return None
|
||||
raw = getattr(patcher, "_patcher", patcher)
|
||||
# PTD installs on the real patcher, not the guarded one. A guard would drop a failing hook and carry on
|
||||
# with the plain sampler on a model whose layers expect PTD's hops, which renders noise; here it ends the run.
|
||||
path = folder_paths.get_full_path("ptd", values["adapter"]) if values["adapter"] else None
|
||||
if path is None:
|
||||
raise control.Unavailable("no PTD adapter file is chosen: put one in ComfyUI/models/ptd, restart, "
|
||||
"and pick it in Settings")
|
||||
native = raw.get_model_object("diffusion_model")
|
||||
adapter = _adapter(path)
|
||||
adapters.verify(adapter, native)
|
||||
if CONTROLLER_KEY in raw.model_options:
|
||||
raise ValueError("This model already carries PTD adapters")
|
||||
controller = Controller(native, adapter)
|
||||
controller.fuse_blocks = KERNELS[values["kernels"]]
|
||||
attach(raw, controller, key, patching.tag)
|
||||
raw.add_wrapper_with_key(WrappersMP.SAMPLER_SAMPLE, key, patching.tag(_swap(controller), key))
|
||||
size = adapter.size / 2**30
|
||||
return (f"{adapter.label}, {adapter.hops} hops, {'int8' if adapter.int8 else 'float'} {size:.1f} GB, "
|
||||
f"{values['kernels']} blocks")
|
||||
|
||||
|
||||
def patcher_has_ptd(patcher):
|
||||
from .controller import CONTROLLER_KEY
|
||||
return CONTROLLER_KEY in patcher.model_options
|
||||
|
||||
PROVIDES = {"modifier": install}
|
||||
@@ -0,0 +1,227 @@
|
||||
"""Reading a PTD adapter file into a bank of per-hop experts.
|
||||
|
||||
Three families load here:
|
||||
|
||||
* EPFL fl2va releases (bf16, or int8 with per-row scales): keys `transformer.<module>.lora_A.<hop>.weight`.
|
||||
* taobao-content-ai Ref2VA releases (bf16, or int8): factors stacked over hops (`lora_down [H, r, in]`)
|
||||
plus a bank of per-interval output heads that is fused per hop.
|
||||
* Any rank: reduced-rank files simply have smaller factors.
|
||||
|
||||
The key mapping, topology checks and head fusion follow ComfyUI-PTD (EPFL VITA Lab) and
|
||||
ComfyUI-MiniMaxH3-PTD (taobao-content-ai), both Apache-2.0, via ComfyUI-PTD-Compact (Nynxz).
|
||||
"""
|
||||
|
||||
from collections import defaultdict
|
||||
from dataclasses import dataclass, field
|
||||
import re
|
||||
|
||||
import torch
|
||||
|
||||
from . import schedule
|
||||
|
||||
HEADS = {"proj_out": "final_layer.video_out", "audio_proj_out": "final_layer.audio_out"}
|
||||
PARTS = ("attn.to_q", "attn.to_k", "attn.to_v", "attn.to_out.0", "ff.net.0.proj", "ff.net.2")
|
||||
_QKV = {"attn.to_q": 0, "attn.to_k": 1, "attn.to_v": 2}
|
||||
_DENSE = {"attn.to_out.0": "attn.out_proj", "ff.net.0.proj": "mlp.fc1", "ff.net.2": "mlp.fc2"}
|
||||
#: Reference tensors some EPFL files carry, to prove the loaded base is the one they were trained on.
|
||||
BUFFERS = {
|
||||
"transformer.rope.inv_freq": "rope.inv_freq",
|
||||
"transformer.proj_out.orig_weight": "final_layer.video_out.weight",
|
||||
"transformer.proj_out.orig_bias": "final_layer.video_out.bias",
|
||||
"transformer.audio_proj_out.orig_weight": "final_layer.audio_out.weight",
|
||||
"transformer.audio_proj_out.orig_bias": "final_layer.audio_out.bias",
|
||||
}
|
||||
_EPFL_KEY = re.compile(r"^transformer\.(?P<module>.+)\.lora_(?P<factor>[AB])\.(?P<hop>0|[1-9][0-9]*)\.weight$")
|
||||
_MODULE = re.compile(r"^(transformer_blocks|token_refiner\.refiner_blocks)\.(\d+)\.(.+)$")
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class Expert:
|
||||
a: torch.Tensor # [r, in] int8 or float
|
||||
b: torch.Tensor # [out, r] int8 or float
|
||||
sa: torch.Tensor | None # [r] float32 row scales of a
|
||||
sb: torch.Tensor | None # [out] float32 row scales of b
|
||||
qkv_index: int | None
|
||||
|
||||
|
||||
@dataclass
|
||||
class Adapter:
|
||||
label: str
|
||||
hops: int
|
||||
grid: tuple # hop times t_0 > ... > 0
|
||||
audio: tuple # per hop, the factor the audio head is scaled by
|
||||
bank: dict # native module -> {hop: (Expert, ...)}
|
||||
heads: list = field(default_factory=list) # per hop {native head: (weight, bias)} (Ref2VA only)
|
||||
int8: bool = False
|
||||
buffers: dict = field(default_factory=dict) # reference tensors of the base (EPFL files)
|
||||
|
||||
@property
|
||||
def size(self):
|
||||
return sum(t.numel() * t.element_size() for by_hop in self.bank.values()
|
||||
for experts in by_hop.values() for e in experts for t in (e.a, e.b))
|
||||
|
||||
|
||||
def native_target(module):
|
||||
"""Diffusers module name -> (native module, qkv index)."""
|
||||
m = _MODULE.fullmatch(module)
|
||||
if m is None:
|
||||
raise ValueError(f"Unknown PTD adapter module: {module}")
|
||||
base = ("blocks." if m[1] == "transformer_blocks" else "token_refiner.blocks.") + m[2]
|
||||
if m[3] in _QKV:
|
||||
return f"{base}.attn.qkv_proj", _QKV[m[3]]
|
||||
if m[3] not in _DENSE:
|
||||
raise ValueError(f"Unknown PTD adapter projection: {module}")
|
||||
return f"{base}.{_DENSE[m[3]]}", None
|
||||
|
||||
|
||||
def _add(bank, name, hop, a, b, sa, sb, qkv):
|
||||
if name.endswith(".mlp.fc1"): # Diffusers stores [up, gate]; native H3 consumes [gate, up]
|
||||
b = torch.cat(b.chunk(2, dim=0)[::-1], dim=0)
|
||||
sb = None if sb is None else torch.cat(sb.chunk(2, dim=0)[::-1], dim=0)
|
||||
bank[name][hop].append(Expert(a.contiguous(), b.contiguous(),
|
||||
None if sa is None else sa.reshape(-1).float(),
|
||||
None if sb is None else sb.reshape(-1).float(), qkv))
|
||||
|
||||
|
||||
def _epfl_bank(tensors, hops):
|
||||
bank = defaultdict(lambda: defaultdict(list))
|
||||
for key in tensors:
|
||||
if key.endswith("_scale") or key in BUFFERS:
|
||||
continue
|
||||
m = _EPFL_KEY.fullmatch(key)
|
||||
if m is None:
|
||||
raise ValueError(f"Not a supported PTD adapter key: {key}")
|
||||
if m["factor"] != "A":
|
||||
continue
|
||||
hop, module = int(m["hop"]), m["module"]
|
||||
if hop >= hops:
|
||||
raise ValueError(f"Adapter hop {hop} is beyond this file's {hops} hops: {key}")
|
||||
b_key = key.replace(".lora_A.", ".lora_B.")
|
||||
a, b = tensors[key], tensors[b_key]
|
||||
if a.ndim != 2 or b.ndim != 2 or a.shape[0] != b.shape[1]:
|
||||
raise ValueError(f"Invalid PTD adapter pair: {key}")
|
||||
if module.endswith(".main"):
|
||||
name, qkv = HEADS[module[:-len(".main")]], None
|
||||
else:
|
||||
name, qkv = native_target(module)
|
||||
_add(bank, name, hop, a, b, tensors.get(key + "_scale"), tensors.get(b_key + "_scale"), qkv)
|
||||
return bank
|
||||
|
||||
|
||||
def _pdd_bank(tensors, hops, scale):
|
||||
bank = defaultdict(lambda: defaultdict(list))
|
||||
for module in sorted({k.rsplit(".lora_", 1)[0] for k in tensors if ".lora_" in k}):
|
||||
name, qkv = native_target(module)
|
||||
down, up = tensors[f"{module}.lora_down"], tensors[f"{module}.lora_up"]
|
||||
sd, su = tensors.get(f"{module}.lora_down_scale"), tensors.get(f"{module}.lora_up_scale")
|
||||
if down.ndim != 3 or up.ndim != 3 or down.shape[0] != hops or up.shape[0] != hops:
|
||||
raise ValueError(f"{module}: expected factors stacked over {hops} hops")
|
||||
for h in range(hops):
|
||||
sb = None if su is None else su[h].float()
|
||||
if scale != 1.0:
|
||||
sb = torch.full((up.shape[1],), scale) if sb is None else sb * scale
|
||||
_add(bank, name, h, down[h], up[h], None if sd is None else sd[h].float(), sb, qkv)
|
||||
return bank
|
||||
|
||||
|
||||
def _pdd_heads(tensors, hops, steps):
|
||||
"""Per hop, the bank of per-interval heads fused with that hop's normalised step sizes."""
|
||||
banks = {}
|
||||
for prefix, native in HEADS.items():
|
||||
if f"{prefix}.weight" not in tensors:
|
||||
raise ValueError(f"PTD file is missing the output head bank {prefix}.weight")
|
||||
w, b = tensors[f"{prefix}.weight"], tensors.get(f"{prefix}.bias")
|
||||
if w.ndim != 3 or w.shape[0] != steps or (b is not None and tuple(b.shape) != (steps, w.shape[1])):
|
||||
raise ValueError(f"Output head bank {prefix} must be ({steps}, out, in)")
|
||||
banks[native] = (w.float(), None if b is None else b.float())
|
||||
heads = []
|
||||
for video_plan, audio_plan in schedule.pdd_plans(hops, steps):
|
||||
fused = {}
|
||||
for name, (w, b) in banks.items():
|
||||
plan = video_plan if name.endswith("video_out") else audio_plan
|
||||
fused[name] = (torch.einsum("n,noi->oi", plan, w),
|
||||
None if b is None else torch.einsum("n,no->o", plan, b))
|
||||
heads.append(fused)
|
||||
return heads
|
||||
|
||||
|
||||
def _expected(native, head_experts):
|
||||
names = set()
|
||||
for prefix, count in (("blocks", len(native.blocks)), ("token_refiner.blocks", len(native.token_refiner.blocks))):
|
||||
for i in range(count):
|
||||
names |= {f"{prefix}.{i}.attn.qkv_proj", f"{prefix}.{i}.attn.out_proj",
|
||||
f"{prefix}.{i}.mlp.fc1", f"{prefix}.{i}.mlp.fc2"}
|
||||
return names | set(HEADS.values()) if head_experts else names
|
||||
|
||||
|
||||
def verify(adapter, native):
|
||||
"""Raise, in plain words, unless `adapter` fits the loaded H3 model `native` layer for layer."""
|
||||
bank, hops = adapter.bank, adapter.hops
|
||||
expected = _expected(native, head_experts=not adapter.heads)
|
||||
if set(bank) != expected:
|
||||
raise ValueError(f"This adapter does not match the loaded model's layers: {len(expected - set(bank))} "
|
||||
f"missing, {len(set(bank) - expected)} unexpected")
|
||||
for name, by_hop in bank.items():
|
||||
module = native.get_submodule(name)
|
||||
if set(by_hop) != set(range(hops)):
|
||||
raise ValueError(f"PTD adapter is missing hops for {name}")
|
||||
for experts in by_hop.values():
|
||||
if len(experts) != (3 if name.endswith("qkv_proj") else 1):
|
||||
raise ValueError(f"PTD adapter has the wrong number of factors for {name}")
|
||||
for e in experts:
|
||||
if e.a.shape[1] != module.in_features or \
|
||||
e.b.shape[0] * (3 if e.qkv_index is not None else 1) != module.out_features:
|
||||
raise ValueError(f"PTD adapter does not fit {name}: it was made for another H3 variant")
|
||||
for name in HEADS.values() if adapter.heads else ():
|
||||
head = native.get_submodule(name)
|
||||
if (head.in_features, head.out_features) != tuple(reversed(adapter.heads[0][name][0].shape)):
|
||||
raise ValueError(f"The output head bank does not fit {name}: it was made for another H3 variant")
|
||||
for key, target in BUFFERS.items():
|
||||
if key not in adapter.buffers:
|
||||
continue
|
||||
parent, _, attr = target.rpartition(".")
|
||||
base = getattr(native.get_submodule(parent), attr)
|
||||
if not isinstance(base, torch.Tensor) or base.device.type == "meta" or \
|
||||
base.dtype not in (torch.float32, torch.float16, torch.bfloat16):
|
||||
continue # a quantised base cannot be compared; the layer checks above still apply
|
||||
value = adapter.buffers[key]
|
||||
if base.shape != value.shape or not torch.equal(base.detach().cpu(), value.to("cpu", base.dtype)):
|
||||
raise ValueError(f"The loaded MiniMax-H3 base is not the one this adapter was trained on ({target})")
|
||||
|
||||
|
||||
def read(path):
|
||||
from safetensors import safe_open
|
||||
with safe_open(str(path), framework="pt", device="cpu") as f:
|
||||
meta = dict(f.metadata() or {})
|
||||
tensors = {k: f.get_tensor(k) for k in f.keys()}
|
||||
for key, value in tensors.items():
|
||||
if value.is_floating_point() and not bool(torch.isfinite(value).all()):
|
||||
raise ValueError(f"Non-finite values in {key}")
|
||||
return tensors, meta
|
||||
|
||||
|
||||
def parse(path, label=None):
|
||||
"""The Adapter in `path`. It knows nothing of any model yet: `verify` it against one before use."""
|
||||
tensors, meta = read(path)
|
||||
buffers = {k: v for k, v in tensors.items() if k in BUFFERS}
|
||||
if "ptd_num_hops" in meta: # Ref2VA release
|
||||
hops = int(meta["ptd_num_hops"])
|
||||
steps = int(meta.get("pdd_num_steps", schedule.PDD_STEPS))
|
||||
if hops < 1 or steps % hops:
|
||||
raise ValueError(f"pdd_num_steps={steps} is not divisible by ptd_num_hops={hops}")
|
||||
rank = int(meta.get("lora_rank", 128))
|
||||
bank = _pdd_bank(tensors, hops, float(meta.get("lora_alpha", rank)) / rank)
|
||||
heads = _pdd_heads(tensors, hops, steps)
|
||||
grid = schedule.pdd_grid(hops, steps)
|
||||
audio = schedule.audio_chords(grid)
|
||||
elif meta.get("ptd_recipe") in schedule.RECIPES: # EPFL release
|
||||
grid = schedule.RECIPES[meta["ptd_recipe"]]
|
||||
hops = len(grid) - 1
|
||||
bank = _epfl_bank(tensors, hops)
|
||||
heads = []
|
||||
audio = schedule.audio_tangents(grid)
|
||||
else:
|
||||
raise ValueError("Not a PTD adapter file: it names neither ptd_recipe nor ptd_num_hops")
|
||||
bank = {name: {h: tuple(e) for h, e in by_hop.items()} for name, by_hop in bank.items()}
|
||||
int8 = any(e.a.dtype == torch.int8 for by_hop in bank.values() for es in by_hop.values() for e in es)
|
||||
return Adapter(label or str(path), hops, grid, audio, bank, heads, int8, buffers)
|
||||
@@ -0,0 +1,384 @@
|
||||
"""Running a PTD adapter on a native MiniMax-H3 model: one expert set per hop, switched by hooks.
|
||||
|
||||
Adapters are low-rank additions to linear outputs (forward hooks), so ordinary LoRAs, which are weight
|
||||
patches, stack with them in either order. The bank stays in system RAM; each layer copies only the
|
||||
current hop's factors to the GPU, and int8 row scales are folded into the small factors.
|
||||
|
||||
Hooks live for one model call and are removed even when the call fails. Block arithmetic follows PTD's
|
||||
reference order (bf16 products rounded one by one, as the experts were trained) unless `fuse_blocks`
|
||||
asks for ComfyUI's fused native blocks.
|
||||
|
||||
Controller, hooks, block maths and the model-call wrapper follow ComfyUI-PTD (EPFL VITA Lab, Apache-2.0),
|
||||
with the int8 factors, in-place residual and DynamicVRAM handling of ComfyUI-PTD-Compact (Nynxz,
|
||||
Apache-2.0) and the output-head fusion of ComfyUI-MiniMaxH3-PTD (taobao-content-ai, Apache-2.0).
|
||||
"""
|
||||
|
||||
from contextlib import ExitStack, contextmanager
|
||||
from contextvars import ContextVar
|
||||
import copy
|
||||
from dataclasses import dataclass
|
||||
import math
|
||||
import threading
|
||||
import weakref
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
from . import schedule
|
||||
|
||||
_FORWARD = ContextVar("funpack_ptd_forward", default=None)
|
||||
_LOCKS = weakref.WeakKeyDictionary()
|
||||
_LOCKS_GUARD = threading.Lock()
|
||||
|
||||
|
||||
def _model_lock(model):
|
||||
with _LOCKS_GUARD:
|
||||
return _LOCKS.setdefault(model, threading.RLock())
|
||||
|
||||
|
||||
def _check_interrupt():
|
||||
from comfy.model_management import throw_exception_if_processing_interrupted
|
||||
throw_exception_if_processing_interrupted()
|
||||
|
||||
|
||||
# --- block arithmetic in PTD's reference order ------------------------------------------------------
|
||||
|
||||
def _modulated(normed, shift, scale, segments):
|
||||
result = torch.empty_like(normed)
|
||||
for start, stop, row in segments:
|
||||
result[start:stop] = normed[start:stop] * (1.0 + scale[row].to(normed.dtype)) + shift[row].to(normed.dtype)
|
||||
return result
|
||||
|
||||
|
||||
def _residual(state, update, gate, segments):
|
||||
"""`state + update * gate`, added in place. The product is rounded on its own (no addcmul) as PTD was
|
||||
trained; adding into the stream keeps ComfyUI's allocation compiler from seeing a block output that
|
||||
outlives its block."""
|
||||
for start, stop, row in segments:
|
||||
state[start:stop].add_(update[start:stop] * gate[row].to(state.dtype))
|
||||
return state
|
||||
|
||||
|
||||
def _block_math(block, x, t_emb, mod_segments, rope_freqs, transformer_options=None, attention=None):
|
||||
options = {} if transformer_options is None else transformer_options
|
||||
shift_a, scale_a, gate_a, shift_m, scale_m, gate_m = block.adaln_proj(t_emb)
|
||||
normed = _modulated(block.norm1(x), shift_a, scale_a, mod_segments)
|
||||
update = (block.attn if attention is None else attention)(normed, rope_freqs=rope_freqs, transformer_options=options)
|
||||
x = _residual(x, update, gate_a, mod_segments)
|
||||
normed = _modulated(block.norm2(x), shift_m, scale_m, mod_segments)
|
||||
return _residual(x, block.mlp(normed), gate_m, mod_segments)
|
||||
|
||||
|
||||
def _split_half_rotation(x, table):
|
||||
cos = table[0, :, 0, :, 0, 0].repeat(1, 2).unsqueeze(1).to(x.dtype)
|
||||
sin = table[0, :, 0, :, 1, 0].repeat(1, 2).unsqueeze(1).to(x.dtype)
|
||||
width = cos.shape[-1]
|
||||
leading = x[..., :width]
|
||||
first, second = leading.chunk(2, dim=-1)
|
||||
# The two products and their sum stay separate: a fused RMS-RoPE would do the rotation in fp32 first.
|
||||
result = leading * cos + torch.cat((-second, first), dim=-1) * sin
|
||||
return torch.cat((result, x[..., width:]), dim=-1)
|
||||
|
||||
|
||||
def _attention_math(attention, x, rope_freqs=None, transformer_options=None):
|
||||
from comfy.ldm.modules.attention import AttentionTensorContainer, optimized_attention
|
||||
options = {} if transformer_options is None else transformer_options
|
||||
streams = attention.qkv_proj(x).chunk(3, dim=-1)
|
||||
q, k, v = [value.reshape(x.shape[0], attention.heads, attention.head_dim) for value in streams]
|
||||
q, k = attention.q_norm(q), attention.k_norm(k)
|
||||
if rope_freqs is not None:
|
||||
q, k = _split_half_rotation(q, rope_freqs), _split_half_rotation(k, rope_freqs)
|
||||
packed = [AttentionTensorContainer(value.transpose(0, 1).unsqueeze(0)) for value in (q, k, v)]
|
||||
result = optimized_attention(*packed, attention.heads, skip_reshape=True,
|
||||
preferred_attention=attention.comfy_attention, transformer_options=options)
|
||||
return attention.out_proj(result.squeeze(0))
|
||||
|
||||
|
||||
def _scope_forward(module, function, frame, originals):
|
||||
original = module.forward
|
||||
originals.append((module, "forward" in vars(module), original))
|
||||
|
||||
def forward(*args, **kwargs):
|
||||
if _FORWARD.get() is not frame:
|
||||
return original(*args, **kwargs)
|
||||
return function(module, *args, **kwargs)
|
||||
|
||||
module.forward = forward
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _without_comfy_compiler():
|
||||
"""Switch off ComfyUI's allocation compiler (and its CUDA graphs) for one call. Used only for a run in
|
||||
which the compiler failed: its recording does not expect the extra tensors PTD's block maths allocates."""
|
||||
try:
|
||||
from comfy.cli_args import args
|
||||
except ImportError:
|
||||
yield
|
||||
return
|
||||
saved = getattr(args, "disable_comfy_compiler", False), getattr(args, "disable_cuda_graphs", False)
|
||||
args.disable_comfy_compiler = args.disable_cuda_graphs = True
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
args.disable_comfy_compiler, args.disable_cuda_graphs = saved
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class Hop:
|
||||
index: int
|
||||
sigma: float
|
||||
|
||||
|
||||
class Controller:
|
||||
"""The adapter bank of one loaded file, bound to one diffusion model. Hop selection is context-local;
|
||||
calls that share a backbone are serialised."""
|
||||
|
||||
def __init__(self, native, adapter):
|
||||
self.native = native
|
||||
self.adapter = adapter
|
||||
self.hops = adapter.hops
|
||||
self.grid = torch.tensor(adapter.grid, dtype=torch.float64)
|
||||
self._lock = _model_lock(native)
|
||||
self._hop = ContextVar(f"funpack_ptd_hop_{id(self)}", default=None)
|
||||
self.fuse_blocks = False
|
||||
self.allocation_compiler = True
|
||||
|
||||
@contextmanager
|
||||
def activate(self, hop, sigma):
|
||||
if isinstance(hop, bool) or not isinstance(hop, int) or not 0 <= hop < self.hops:
|
||||
raise ValueError(f"PTD hop must be in [0, {self.hops})")
|
||||
expected = float(self.grid[hop])
|
||||
if abs(float(sigma) - expected) > 2e-6:
|
||||
raise ValueError(f"Hop {hop} must run at its trained time {expected:.6f}, not {float(sigma):.6f}")
|
||||
token = self._hop.set(Hop(hop, expected))
|
||||
try:
|
||||
yield self
|
||||
finally:
|
||||
self._hop.reset(token)
|
||||
|
||||
def current(self):
|
||||
hop = self._hop.get()
|
||||
if hop is None:
|
||||
raise RuntimeError("PTD adapters are on this model but no PTD hop is running: the model was called "
|
||||
"outside the PTD sampler")
|
||||
return hop
|
||||
|
||||
def _linear_hook(self, frame, experts):
|
||||
# q/k/v share one input: their A factors run as one matmul. Each expert's delta is added into the
|
||||
# fresh linear output in place; int8 row scales are folded into the small factors.
|
||||
qkv = sorted((e for e in experts if e.qkv_index is not None), key=lambda e: e.qkv_index)
|
||||
plain = [e for e in experts if e.qkv_index is None]
|
||||
ranks = [e.a.shape[0] for e in qkv]
|
||||
|
||||
def factors(e, dev, dt):
|
||||
a = e.a.to(dev, non_blocking=True).to(dt)
|
||||
b = e.b.to(dev, non_blocking=True).to(dt)
|
||||
if e.sa is not None:
|
||||
a = a * e.sa.to(dev, dt)[:, None]
|
||||
if e.sb is not None:
|
||||
b = b * e.sb.to(dev, dt)[:, None]
|
||||
return a, b
|
||||
|
||||
def hook(module, args, kwargs, output):
|
||||
if _FORWARD.get() is not frame:
|
||||
return output
|
||||
_check_interrupt()
|
||||
x = args[0] if args else kwargs.get("input")
|
||||
dev, dt = x.device, output.dtype # the layer's compute dtype
|
||||
x = x.to(dt)
|
||||
for e in plain:
|
||||
a, b = factors(e, dev, dt)
|
||||
output.add_(F.linear(F.linear(x, a), b))
|
||||
if qkv:
|
||||
pairs = [factors(e, dev, dt) for e in qkv]
|
||||
hs = F.linear(x, torch.cat([a for a, _ in pairs])).split(ranks, dim=-1)
|
||||
width = output.shape[-1] // 3
|
||||
for e, (_, b), h in zip(qkv, pairs, hs):
|
||||
output[..., e.qkv_index * width:(e.qkv_index + 1) * width].add_(F.linear(h, b))
|
||||
return output
|
||||
return hook
|
||||
|
||||
@contextmanager
|
||||
def forward_hooks(self, hop):
|
||||
"""Scope the adapters, the PTD clocks and (unless fuse_blocks) PTD's block maths to one model call."""
|
||||
native = self.native
|
||||
with self._lock, ExitStack() as stack:
|
||||
if not self.allocation_compiler:
|
||||
stack.enter_context(_without_comfy_compiler())
|
||||
frame = object()
|
||||
token = _FORWARD.set(frame)
|
||||
handles, forced, forwards = [], [], []
|
||||
try:
|
||||
if not self.fuse_blocks:
|
||||
for block in native.blocks:
|
||||
if all(hasattr(block, n) for n in ("norm1", "norm2", "adaln_proj", "attn", "mlp")):
|
||||
_scope_forward(block, _block_math, frame, forwards)
|
||||
_scope_forward(block.attn, _attention_math, frame, forwards)
|
||||
self._time_hooks(frame, hop, handles)
|
||||
heads = self.adapter.heads[hop.index] if self.adapter.heads else {}
|
||||
for name, (weight, bias) in heads.items():
|
||||
parent, _, attr = name.rpartition(".")
|
||||
handles.append(getattr(native.get_submodule(parent), attr).register_forward_hook(
|
||||
self._fused_head(frame, weight, bias)))
|
||||
for name, by_hop in self.adapter.bank.items():
|
||||
module = native.get_submodule(name)
|
||||
# Native SwiGLU may fuse activation and fc2, bypassing Module.__call__ and its hooks;
|
||||
# ComfyUI's public flag routes the call through the plain linear.
|
||||
if hasattr(module, "comfy_force_forward") or name.endswith(".mlp.fc2"):
|
||||
forced.append((module, "comfy_force_forward" in vars(module),
|
||||
getattr(module, "comfy_force_forward", None)))
|
||||
module.comfy_force_forward = True
|
||||
handles.append(module.register_forward_hook(
|
||||
self._linear_hook(frame, self.adapter.bank[name][hop.index]), with_kwargs=True))
|
||||
factor = self.adapter.audio[hop.index]
|
||||
|
||||
def audio_clock(module, args, output):
|
||||
if _FORWARD.get() is not frame:
|
||||
return output
|
||||
# Scaled before native H3 casts to the latent dtype, like PTD's fp32 output heads.
|
||||
return output * torch.as_tensor(factor, dtype=output.dtype, device=output.device)
|
||||
|
||||
handles.append(native.final_layer.audio_out.register_forward_hook(audio_clock))
|
||||
_check_interrupt()
|
||||
yield
|
||||
_check_interrupt()
|
||||
finally:
|
||||
for handle in reversed(handles):
|
||||
handle.remove()
|
||||
for module, had, value in reversed(forced):
|
||||
if had:
|
||||
module.comfy_force_forward = value
|
||||
else:
|
||||
del module.comfy_force_forward
|
||||
for module, had, original in reversed(forwards):
|
||||
if had:
|
||||
module.forward = original
|
||||
else:
|
||||
del module.forward
|
||||
_FORWARD.reset(token)
|
||||
|
||||
@staticmethod
|
||||
def _fused_head(frame, weight, bias):
|
||||
def hook(module, args, output):
|
||||
if _FORWARD.get() is not frame:
|
||||
return output
|
||||
x = args[0]
|
||||
return F.linear(x.to(output.dtype), weight.to(x.device, output.dtype),
|
||||
None if bias is None else bias.to(x.device, output.dtype))
|
||||
return hook
|
||||
|
||||
def _time_hooks(self, frame, hop, handles):
|
||||
native = self.native
|
||||
if not hasattr(native, "time_embedder"): # pruned (AdaLN-curve) base: no time embedder to steer
|
||||
return
|
||||
remembered = {}
|
||||
|
||||
def exact_times(module, args):
|
||||
if _FORWARD.get() is not frame:
|
||||
return args
|
||||
times = args[0]
|
||||
if not isinstance(times, torch.Tensor) or times.ndim != 1:
|
||||
raise ValueError("PTD expects a vector of native H3 time labels")
|
||||
video_time = 1.0 - hop.sigma
|
||||
candidates = torch.tensor([video_time, 1.0 - schedule.audio_sigma(hop.sigma), max(video_time, .999)],
|
||||
device=times.device, dtype=torch.float32).unique()
|
||||
distance, indices = (times[:, None] - candidates[None]).abs().min(dim=1)
|
||||
if not bool((distance <= 2e-6).all()):
|
||||
raise ValueError("Native H3 time labels do not match the active PTD clocks")
|
||||
# Native H3 rounds sigma to fp32 first; PTD subtracts in float64. Keep the table order.
|
||||
return (candidates[indices],) + args[1:]
|
||||
|
||||
def remember(module, args, output):
|
||||
if _FORWARD.get() is frame:
|
||||
remembered["fp32"] = output
|
||||
|
||||
def precise_adaln_input(module, args):
|
||||
if _FORWARD.get() is not frame:
|
||||
return args
|
||||
if "fp32" not in remembered:
|
||||
raise RuntimeError("PTD AdaLN ran before the time embedder")
|
||||
# Diffusers applies SiLU to the fp32 time embedding before casting; native H3 casts first.
|
||||
return (F.silu(remembered["fp32"]).to(args[0].dtype),) + args[1:]
|
||||
|
||||
handles.append(native.time_embedder.register_forward_pre_hook(exact_times))
|
||||
handles.append(native.time_embedder.register_forward_hook(remember))
|
||||
for layer in (*native.blocks, native.final_layer):
|
||||
adaln = getattr(layer, "adaln_proj", None)
|
||||
if adaln is not None and getattr(adaln, "apply_silu", True):
|
||||
handles.append(adaln.linear.register_forward_pre_hook(precise_adaln_input))
|
||||
|
||||
|
||||
class ModelCall:
|
||||
"""The model-function wrapper: PTD carries physical latents, so the native shifts are neutralised for the
|
||||
call and the adapters are switched on for it."""
|
||||
|
||||
def __init__(self, controller, previous=None):
|
||||
self.controller = controller
|
||||
self.previous = previous
|
||||
|
||||
def __call__(self, model_function, kwargs):
|
||||
hop = self.controller.current()
|
||||
sigma = kwargs.get("timestep")
|
||||
if not isinstance(sigma, torch.Tensor) or sigma.numel() == 0:
|
||||
raise ValueError("PTD needs a tensor sigma from the sampler")
|
||||
if not bool(torch.allclose(sigma, torch.full_like(sigma, hop.sigma), rtol=0.0, atol=2e-6)):
|
||||
raise ValueError("The model sigma does not match the active PTD hop")
|
||||
forwarded = dict(kwargs)
|
||||
conditions = dict(kwargs.get("c", {}))
|
||||
options = dict(conditions.get("transformer_options", {}))
|
||||
options["minimax_h3_sigma_shift_video"] = 1.0
|
||||
options["minimax_h3_sigma_shift_audio"] = schedule.native_audio_shift(hop.sigma)
|
||||
conditions["transformer_options"] = options
|
||||
payload = dict(conditions.get("minimax_payload", {}))
|
||||
payload["audio_scale"] = 1.0
|
||||
conditions["minimax_payload"] = payload
|
||||
forwarded["c"] = conditions
|
||||
with self.controller.forward_hooks(hop):
|
||||
if self.previous is not None:
|
||||
return self.previous(model_function, forwarded)
|
||||
return model_function(forwarded["input"], forwarded["timestep"], **conditions)
|
||||
|
||||
def to(self, device):
|
||||
if self.previous is not None and hasattr(self.previous, "to"):
|
||||
self.previous = self.previous.to(device)
|
||||
return self
|
||||
|
||||
def cleanup(self, **kwargs):
|
||||
if self.previous is not None and hasattr(self.previous, "cleanup"):
|
||||
self.previous.cleanup(**kwargs)
|
||||
|
||||
def models(self):
|
||||
if self.previous is not None and hasattr(self.previous, "models"):
|
||||
return self.previous.models()
|
||||
return []
|
||||
|
||||
|
||||
def attach(patcher, controller, key, tag):
|
||||
"""Put `controller` on `patcher` (a clone the caller owns): raw text embeddings so each hop's refiner
|
||||
adapters run, PTD's audio shift, and the model-call wrapper. Everything is tagged `key` for removal."""
|
||||
native = controller.native
|
||||
if not all(hasattr(native, n) for n in ("blocks", "token_refiner", "condition_proj", "final_layer",
|
||||
"sigma_shift_video", "sigma_shift_audio")):
|
||||
raise ValueError("PTD needs a native MiniMax-H3 model")
|
||||
sampling = copy.deepcopy(patcher.get_model_object("model_sampling"))
|
||||
if not hasattr(sampling, "audio_shift") or not hasattr(sampling, "shift"):
|
||||
raise ValueError("PTD needs the native MiniMax-H3 audio/video flow sampling")
|
||||
sampling.audio_shift = sampling.shift
|
||||
patcher.add_object_patch("model_sampling", tag(sampling, key))
|
||||
|
||||
text_dim = native.condition_proj.in_features
|
||||
|
||||
def raw_text(text_states):
|
||||
if text_states.shape[-1] != text_dim:
|
||||
raise ValueError("PTD needs raw MiniMax-H3 text conditioning, not pre-refined embeddings")
|
||||
return text_states
|
||||
|
||||
patcher.add_object_patch("diffusion_model.preprocess_text_embeds", tag(raw_text, key))
|
||||
previous = patcher.model_options.get("model_function_wrapper")
|
||||
wrapper = tag(ModelCall(controller, previous), key)
|
||||
wrapper._funpack_prev = previous
|
||||
patcher.set_model_unet_function_wrapper(wrapper)
|
||||
patcher.model_options[CONTROLLER_KEY] = tag(controller, key)
|
||||
|
||||
|
||||
CONTROLLER_KEY = "funpack_ptd_controller"
|
||||
@@ -0,0 +1,135 @@
|
||||
"""PTD's fixed-hop sampler: one trained expert set per hop, Euler on the shared-clock velocity, state in fp32.
|
||||
|
||||
The recurrence follows ComfyUI-PTD (EPFL VITA Lab, Apache-2.0); the console bar, VRAM housekeeping between
|
||||
hops and the retry without ComfyUI's allocation compiler follow ComfyUI-PTD-Compact (Nynxz, Apache-2.0).
|
||||
"""
|
||||
|
||||
import time
|
||||
|
||||
import torch
|
||||
|
||||
from ..._core import log
|
||||
from . import schedule
|
||||
|
||||
|
||||
def _streams(value):
|
||||
if getattr(value, "is_nested", False):
|
||||
streams = tuple(value.unbind())
|
||||
if len(streams) != 2:
|
||||
raise ValueError("PTD expects exactly a picture and a sound stream")
|
||||
return streams
|
||||
return (value,)
|
||||
|
||||
|
||||
def _map(value, op):
|
||||
result = [op(s) for s in _streams(value)]
|
||||
return type(value)(result) if getattr(value, "is_nested", False) else result[0]
|
||||
|
||||
|
||||
def _pair(left, right, op):
|
||||
a, b = _streams(left), _streams(right)
|
||||
if len(a) != len(b) or any(x.shape != y.shape for x, y in zip(a, b)):
|
||||
raise ValueError("The model output must keep the latent's streams and shapes")
|
||||
result = [op(x, y) for x, y in zip(a, b)]
|
||||
return type(left)(result) if getattr(left, "is_nested", False) else result[0]
|
||||
|
||||
|
||||
def _refuse_unsupported(model, extra_args):
|
||||
if extra_args.get("denoise_mask") is not None:
|
||||
raise ValueError("PTD cannot work on a masked or partly denoised latent: it always starts from pure noise")
|
||||
start = getattr(model, "latent_image", None)
|
||||
if start is not None and any(bool(torch.count_nonzero(s)) for s in _streams(start)):
|
||||
raise ValueError("PTD starts from an empty latent; give the picture as a reference instead of as the "
|
||||
"starting latent (no second pass, no image-to-image)")
|
||||
|
||||
|
||||
def _between_hops(model, hop):
|
||||
"""Give PyTorch's cached activation memory back after a hop and log where the GPU memory is.
|
||||
|
||||
A hop's activations are freed when it ends but stay cached, and with DynamicVRAM the base model's
|
||||
weights are evicted to make room; raising the watermark again lets them fault back in before the
|
||||
next hop. Housekeeping only: never stops a run."""
|
||||
try:
|
||||
if not torch.cuda.is_available():
|
||||
return
|
||||
before = torch.cuda.memory_reserved() / 2**30
|
||||
torch.cuda.empty_cache()
|
||||
patcher = getattr(getattr(model, "inner_model", None), "model_patcher", None)
|
||||
vbar = patcher._vbar_get() if hasattr(patcher, "_vbar_get") else None
|
||||
if vbar is not None:
|
||||
import comfy.model_management
|
||||
comfy.model_management.reset_cast_buffers()
|
||||
vbar.set_watermark(vbar.max_size)
|
||||
free, total = torch.cuda.mem_get_info()
|
||||
resident = patcher.loaded_size() / 2**30 if patcher is not None else float("nan")
|
||||
log.info("FunPack PTD", f"hop {hop} memory: model resident {resident:.1f} GB; cache {before:.1f} -> "
|
||||
f"{torch.cuda.memory_reserved() / 2**30:.1f} GB; card {(total - free) / 2**30:.1f}/{total / 2**30:.1f} GB used")
|
||||
except Exception as exc: # noqa: BLE001
|
||||
log.once("ptd:memory", log.WARNING, "FunPack PTD", f"could not release cached memory between hops: {exc}")
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def _run(model, x, extra_args, controller, step):
|
||||
state = _map(model.noise, lambda s: (s.to(torch.float64) * schedule.T_MAX).to(torch.float32)) \
|
||||
if getattr(model, "noise", None) is not None else _map(x, lambda s: s.float())
|
||||
_pair(state, x, lambda a, b: a)
|
||||
first = _streams(state)[0]
|
||||
grid = controller.grid.to(first.device)
|
||||
for hop in range(controller.hops):
|
||||
time_now = float(grid[hop])
|
||||
with controller.activate(hop, time_now):
|
||||
denoised = model(state, grid[hop].expand(first.shape[0]), **extra_args)
|
||||
velocity = _pair(state, denoised, lambda s, d: (s.float() - d.float()) / time_now)
|
||||
step(hop, {"x": state, "i": hop, "sigma": grid[hop], "sigma_hat": grid[hop], "denoised": denoised})
|
||||
delta = (grid[hop + 1] - grid[hop]).float()
|
||||
state = _pair(state, velocity, lambda s, v: s.float() + delta * v.float())
|
||||
return state
|
||||
|
||||
|
||||
def sample_ptd(model, x, sigmas, extra_args=None, callback=None, disable=None, *, controller):
|
||||
"""A KSAMPLER function. `sigmas` must be the adapter's own hop times."""
|
||||
from tqdm.auto import tqdm
|
||||
expected = controller.grid
|
||||
if sigmas.shape != expected.shape or not torch.allclose(sigmas.detach().cpu().double(), expected, atol=2e-6):
|
||||
raise ValueError(f"PTD runs only the {controller.hops} hops it was trained on; its own times are "
|
||||
f"{', '.join(f'{float(t):.3f}' for t in expected)}")
|
||||
extra_args = {} if extra_args is None else dict(extra_args)
|
||||
_refuse_unsupported(model, extra_args)
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.empty_cache() # whatever the last node cached is VRAM the allocation compiler cannot get
|
||||
bar = tqdm(total=controller.hops, disable=disable)
|
||||
done = [0]
|
||||
|
||||
def step(hop, record):
|
||||
done[0] += 1
|
||||
bar.update(1)
|
||||
_between_hops(model, hop)
|
||||
if callback is not None:
|
||||
callback(record)
|
||||
|
||||
started = time.perf_counter()
|
||||
try:
|
||||
try:
|
||||
out = _run(model, x, extra_args, controller, step)
|
||||
except RuntimeError as exc:
|
||||
if "aimdo memory compile error" not in str(exc) or not controller.allocation_compiler:
|
||||
raise
|
||||
# A known intermittent ComfyUI failure when VRAM is nearly full. Drop the half-recorded graph and
|
||||
# sample this run again without the compiler; the hops are deterministic, so nothing changes.
|
||||
log.warning("FunPack PTD", f"ComfyUI's allocation compiler failed at hop {done[0]} ({exc}); "
|
||||
"sampling this run again without it")
|
||||
import comfy.model_prefetch
|
||||
comfy.model_prefetch.cleanup_prefetch_queues()
|
||||
controller.allocation_compiler = False
|
||||
bar.reset()
|
||||
done[0] = 0
|
||||
started = time.perf_counter()
|
||||
try:
|
||||
out = _run(model, x, extra_args, controller, step)
|
||||
finally:
|
||||
controller.allocation_compiler = True
|
||||
finally:
|
||||
bar.close()
|
||||
seconds = time.perf_counter() - started
|
||||
log.info("FunPack PTD", f"{controller.hops} hops in {seconds:.1f} s ({seconds / controller.hops:.2f} s/hop)")
|
||||
return out
|
||||
@@ -0,0 +1,92 @@
|
||||
"""Hop times and audio clocks of the PTD releases.
|
||||
|
||||
The EPFL grids and the shared video/audio clock follow ComfyUI-PTD (EPFL VITA Lab, Apache-2.0); the
|
||||
32-interval PDD grid, head plans and chord correction follow ComfyUI-MiniMaxH3-PTD (taobao-content-ai,
|
||||
Apache-2.0), both via ComfyUI-PTD-Compact (Nynxz). See the README Thank you section.
|
||||
"""
|
||||
|
||||
from itertools import pairwise
|
||||
|
||||
import torch
|
||||
|
||||
T_MAX = 0.999 # upper end of the shared clock
|
||||
AUDIO_RELATIVE_SHIFT = 0.25
|
||||
|
||||
#: EPFL releases: hop times t_0 > ... > t_N = 0.
|
||||
RECIPES = {
|
||||
"h3_4step": (0.999, 0.947, 0.857, 0.667, 0.0),
|
||||
"h3_6step": (0.999, 0.968, 0.923, 0.857, 0.750, 0.545, 0.0),
|
||||
}
|
||||
|
||||
PDD_STEPS = 32 # intervals of the Ref2VA releases' own grid
|
||||
VIDEO_SHIFT = 12.0
|
||||
AUDIO_SHIFT = 3.0
|
||||
|
||||
|
||||
def _time(value):
|
||||
value = float(value)
|
||||
if not 0.0 <= value <= T_MAX:
|
||||
raise ValueError(f"PTD time must lie in [0, {T_MAX}], got {value}")
|
||||
return value
|
||||
|
||||
|
||||
def audio_sigma(time):
|
||||
"""Audio's noise level on the shared clock."""
|
||||
time = _time(time)
|
||||
if time == T_MAX:
|
||||
return T_MAX
|
||||
return AUDIO_RELATIVE_SHIFT * time * T_MAX / (T_MAX + (AUDIO_RELATIVE_SHIFT - 1.0) * time)
|
||||
|
||||
|
||||
def audio_sigma_prime(time):
|
||||
"""d(audio_sigma)/dt."""
|
||||
time = _time(time)
|
||||
return AUDIO_RELATIVE_SHIFT * T_MAX ** 2 / (T_MAX + (AUDIO_RELATIVE_SHIFT - 1.0) * time) ** 2
|
||||
|
||||
|
||||
def native_audio_shift(time):
|
||||
"""The T=1 shift that gives native H3 PTD's audio sigma at this time (video shift 1)."""
|
||||
time = _time(time)
|
||||
if time == 0.0:
|
||||
return AUDIO_RELATIVE_SHIFT
|
||||
if time == T_MAX:
|
||||
return 1.0
|
||||
sigma = audio_sigma(time)
|
||||
return sigma * (1.0 - time) / (time * (1.0 - sigma))
|
||||
|
||||
|
||||
def _shifted(shift, sigma):
|
||||
return shift * sigma / (1 + (shift - 1) * sigma)
|
||||
|
||||
|
||||
def pdd_grid(hops, steps=PDD_STEPS, shift=VIDEO_SHIFT):
|
||||
"""Hop times of a Ref2VA file: the video block boundaries of its 32-interval grid, then 0."""
|
||||
size = steps // hops
|
||||
noise = _shifted(float(shift), torch.linspace(1.0, 0.0, steps + 1, dtype=torch.float64))
|
||||
return tuple(T_MAX * float(noise[h * size]) for h in range(hops)) + (0.0,)
|
||||
|
||||
|
||||
def pdd_plans(hops, steps=PDD_STEPS, video_shift=VIDEO_SHIFT, audio_shift=AUDIO_SHIFT):
|
||||
"""Per hop (video, audio) weights over the per-interval heads: the block's normalised step sizes."""
|
||||
size = steps // hops
|
||||
plans = []
|
||||
for shift in (video_shift, audio_shift):
|
||||
sizes = (1.0 - _shifted(float(shift), torch.linspace(1.0, 0.0, steps + 1, dtype=torch.float64))).diff()
|
||||
rows = []
|
||||
for h in range(hops):
|
||||
plan = torch.zeros(steps, dtype=torch.float64)
|
||||
block = sizes[h * size:(h + 1) * size]
|
||||
plan[h * size:(h + 1) * size] = block / block.sum()
|
||||
rows.append(plan.float())
|
||||
plans.append(rows)
|
||||
return list(zip(*plans))
|
||||
|
||||
|
||||
def audio_chords(grid):
|
||||
"""Per hop, the audio clock's mean slope. The Ref2VA audio head is the block-mean velocity."""
|
||||
return tuple((audio_sigma(a) - audio_sigma(b)) / (a - b) for a, b in pairwise(grid))
|
||||
|
||||
|
||||
def audio_tangents(grid):
|
||||
"""Per hop, the audio clock's slope at the hop start. The EPFL audio head is an instant velocity."""
|
||||
return tuple(audio_sigma_prime(a) for a in grid[:-1])
|
||||
@@ -0,0 +1,104 @@
|
||||
"""A tiny real native MiniMax-H3, and PTD adapter files in the three layouts, for the PTD tests.
|
||||
|
||||
The model builder and file layouts follow ComfyUI-PTD-Compact's tests (Nynxz, Apache-2.0).
|
||||
"""
|
||||
|
||||
import json
|
||||
|
||||
import torch
|
||||
from safetensors.torch import save_file
|
||||
|
||||
DIMS = dict(hidden_size=64, num_layers=2, token_refiner_num_layers=1, num_attention_heads=2, attention_head_dim=32,
|
||||
ffn_hidden_size=48, text_dim=40, timestep_input_dim=16, time_embed_hidden_size=32, time_embed_dim=16,
|
||||
rope_inv_freq_len=4)
|
||||
PARTS = {"attn.to_q": ("attn.qkv_proj", 3), "attn.to_k": ("attn.qkv_proj", 3), "attn.to_v": ("attn.qkv_proj", 3),
|
||||
"attn.to_out.0": ("attn.out_proj", 1), "ff.net.0.proj": ("mlp.fc1", 1), "ff.net.2": ("mlp.fc2", 1)}
|
||||
BLOCKS = (("transformer_blocks", "blocks", lambda n: len(n.blocks)),
|
||||
("token_refiner.refiner_blocks", "token_refiner.blocks", lambda n: len(n.token_refiner.blocks)))
|
||||
|
||||
|
||||
def native_model(curve=False, **override):
|
||||
"""A ModelPatcher around a real MiniMaxH3 (random weights, CPU). curve=True is the pruned layout."""
|
||||
import comfy.model_patcher
|
||||
import comfy.supported_models
|
||||
extra = {"adaln_curve_grid": 9} if curve else {}
|
||||
config = comfy.supported_models.MiniMaxH3(dict(image_model="minimax_h3", **{**DIMS, **override}, **extra))
|
||||
config.set_inference_dtype(torch.float32, None)
|
||||
model = config.get_model({})
|
||||
torch.manual_seed(0)
|
||||
with torch.no_grad():
|
||||
for p in model.diffusion_model.parameters():
|
||||
p.normal_(0, 0.05)
|
||||
dm = model.diffusion_model
|
||||
n = dm.rope.inv_freq.numel()
|
||||
dm.rope.inv_freq.copy_(1.0 / 10000 ** (torch.arange(n, dtype=torch.float32) / n))
|
||||
if curve:
|
||||
dm.adaln_t_table.normal_(0, 0.5)
|
||||
cpu = torch.device("cpu")
|
||||
return comfy.model_patcher.ModelPatcher(model, cpu, cpu)
|
||||
|
||||
|
||||
def _int8(x):
|
||||
s = x.abs().amax(-1, keepdim=True).clamp_min(1e-12) / 127
|
||||
return (x / s).round().clamp(-127, 127).to(torch.int8), s.squeeze(-1)
|
||||
|
||||
|
||||
def epfl_file(path, native, hops=4, int8=False, zero_a=False, ranks=(3, 4, 5)):
|
||||
"""An EPFL-layout file; returns {(native module, part, hop): dense delta}."""
|
||||
g = torch.Generator().manual_seed(1)
|
||||
tensors, dense = {}, {}
|
||||
for src, dst, count in BLOCKS:
|
||||
for i in range(count(native)):
|
||||
for part, (module, div) in PARTS.items():
|
||||
layer = native.get_submodule(f"{dst}.{i}.{module}")
|
||||
for hop in range(hops):
|
||||
r = ranks[(hop + len(part)) % len(ranks)]
|
||||
A = torch.randn(r, layer.in_features, generator=g) * 0.1 * (0 if zero_a else 1)
|
||||
B = torch.randn(layer.out_features // div, r, generator=g) * 0.1
|
||||
key = f"transformer.{src}.{i}.{part}"
|
||||
if int8:
|
||||
qa, sa = _int8(A)
|
||||
qb, sb = _int8(B)
|
||||
tensors[f"{key}.lora_A.{hop}.weight"], tensors[f"{key}.lora_A.{hop}.weight_scale"] = qa, sa[:, None]
|
||||
tensors[f"{key}.lora_B.{hop}.weight"], tensors[f"{key}.lora_B.{hop}.weight_scale"] = qb, sb[:, None]
|
||||
A, B = qa.float() * sa[:, None], qb.float() * sb[:, None]
|
||||
else:
|
||||
tensors[f"{key}.lora_A.{hop}.weight"], tensors[f"{key}.lora_B.{hop}.weight"] = A, B
|
||||
dense[(f"{dst}.{i}.{module}", part, hop)] = B @ A
|
||||
for head, native_head in (("proj_out", "video_out"), ("audio_proj_out", "audio_out")):
|
||||
layer = native.final_layer.get_submodule(native_head)
|
||||
for hop in range(hops):
|
||||
A = torch.randn(layer.out_features, layer.in_features, generator=g) * 0.01 * (0 if zero_a else 1)
|
||||
tensors[f"transformer.{head}.main.lora_A.{hop}.weight"] = A
|
||||
tensors[f"transformer.{head}.main.lora_B.{hop}.weight"] = torch.eye(layer.out_features)
|
||||
dense[(f"final_layer.{native_head}", "head", hop)] = A
|
||||
meta = {"ptd_format": "comfyui-native-h3-compact" if int8 else "comfyui-native-h3",
|
||||
"ptd_recipe": f"h3_{hops}step"}
|
||||
save_file(tensors, str(path), metadata=meta)
|
||||
return dense
|
||||
|
||||
|
||||
def pdd_file(path, native, hops=4, steps=32, rank=8, int8=False, zero=False):
|
||||
"""A Ref2VA-layout file: factors stacked over hops, plus a bank of per-interval heads."""
|
||||
g = torch.Generator().manual_seed(1)
|
||||
t = {}
|
||||
for src, dst, count in BLOCKS:
|
||||
for i in range(count(native)):
|
||||
for part, (module, div) in PARTS.items():
|
||||
layer = native.get_submodule(f"{dst}.{i}.{module}")
|
||||
down = torch.randn(hops, rank, layer.in_features, generator=g) * 0.1 * (0 if zero else 1)
|
||||
up = torch.randn(hops, layer.out_features // div, rank, generator=g) * 0.1
|
||||
key = f"{src}.{i}.{part}"
|
||||
if int8:
|
||||
(t[f"{key}.lora_down"], t[f"{key}.lora_down_scale"]) = _int8(down)
|
||||
(t[f"{key}.lora_up"], t[f"{key}.lora_up_scale"]) = _int8(up)
|
||||
else:
|
||||
t[f"{key}.lora_down"], t[f"{key}.lora_up"] = down.bfloat16(), up.bfloat16()
|
||||
for prefix, head in (("proj_out", native.final_layer.video_out), ("audio_proj_out", native.final_layer.audio_out)):
|
||||
w = head.weight.detach().float()
|
||||
t[f"{prefix}.weight"] = (w[None] + torch.randn(steps, *w.shape, generator=g) * 0.05).bfloat16()
|
||||
t[f"{prefix}.bias"] = (torch.randn(steps, w.shape[0], generator=g) * 0.05).bfloat16()
|
||||
meta = {"ptd_num_hops": str(hops), "pdd_num_steps": str(steps), "lora_rank": str(rank), "lora_alpha": f"{rank}.0"}
|
||||
if int8:
|
||||
meta["ptd_compact"] = json.dumps({"schema": 1})
|
||||
save_file(t, str(path), metadata=meta)
|
||||
@@ -0,0 +1,302 @@
|
||||
"""PTD on a tiny real MiniMax-H3: all three file layouts, the hop switching, the sampler swap and clean removal.
|
||||
|
||||
Structure follows ComfyUI-PTD-Compact's tests (Nynxz, Apache-2.0). They prove the wiring and the maths of
|
||||
our side on random weights, nothing about picture quality on the real model.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
from . import _tiny
|
||||
|
||||
HOP_TIMES = (0.999, 0.947, 0.857, 0.667)
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _needs_comfy(comfyui):
|
||||
"""Imports comfy."""
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def ptd():
|
||||
import importlib
|
||||
return importlib.import_module("modules.sampling.ptd")
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def folder(tmp_path):
|
||||
import folder_paths
|
||||
folder_paths.add_model_folder_path("ptd", str(tmp_path), is_default=True)
|
||||
return tmp_path
|
||||
|
||||
|
||||
def _guarded(patcher):
|
||||
from modules._core import patching
|
||||
return patching.GuardedPatcher(patcher, "funpack.ptd", patching.Dropped())
|
||||
|
||||
|
||||
def _install(ptd, patcher, name, kernels="exact"):
|
||||
note = ptd.install(_guarded(patcher), {"enabled": True, "adapter": name, "kernels": kernels}, "funpack.ptd")
|
||||
return patcher.model_options["funpack_ptd_controller"], note
|
||||
|
||||
|
||||
def _clone(base):
|
||||
return base.clone()
|
||||
|
||||
|
||||
def _sample(model, sampler=None, sigmas=None):
|
||||
import comfy.nested_tensor
|
||||
import comfy.samplers
|
||||
guider = comfy.samplers.CFGGuider(model)
|
||||
context = torch.randn(1, 6, _tiny.DIMS["text_dim"], generator=torch.Generator().manual_seed(3))
|
||||
guider.set_conds([[context, {}]], [[context, {}]])
|
||||
guider.set_cfg(1.0)
|
||||
video, audio = torch.zeros(1, 24, 2, 4, 4), torch.zeros(1, 32, 2, 3)
|
||||
noise = comfy.nested_tensor.NestedTensor([torch.randn(video.shape, generator=torch.Generator().manual_seed(4)),
|
||||
torch.randn(audio.shape, generator=torch.Generator().manual_seed(5))])
|
||||
sampler = sampler or comfy.samplers.sampler_object("euler")
|
||||
sigmas = sigmas if sigmas is not None else torch.linspace(1.0, 0.0, 21)
|
||||
try:
|
||||
out = guider.sample(noise, comfy.nested_tensor.NestedTensor([video, audio]), sampler, sigmas,
|
||||
disable_pbar=True, seed=7)
|
||||
finally:
|
||||
import comfy.model_management
|
||||
comfy.model_management.unload_all_models()
|
||||
return [t.clone() for t in out.unbind()]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("int8", [False, True], ids=["float", "int8"])
|
||||
@pytest.mark.parametrize("curve", [False, True], ids=["full", "pruned"])
|
||||
def test_each_hop_adds_its_own_expert_to_the_layer(ptd, folder, int8, curve):
|
||||
base = _tiny.native_model(curve)
|
||||
native = base.get_model_object("diffusion_model")
|
||||
weights = {n: p.detach().clone() for n, p in native.named_parameters()}
|
||||
dense = _tiny.epfl_file(folder / "a.safetensors", native, int8=int8)
|
||||
controller, _ = _install(ptd, _clone(base), "a.safetensors")
|
||||
for hop, sigma in enumerate(HOP_TIMES):
|
||||
for name in ("blocks.0.attn.qkv_proj", "blocks.1.mlp.fc1", "blocks.1.mlp.fc2", "blocks.0.attn.out_proj",
|
||||
"token_refiner.blocks.0.attn.qkv_proj"):
|
||||
layer = native.get_submodule(name)
|
||||
inp = torch.randn(5, layer.in_features)
|
||||
parts = [(p, d) for (n, p, h), d in dense.items() if n == name and h == hop]
|
||||
if name.endswith("qkv_proj"):
|
||||
order = {"attn.to_q": 0, "attn.to_k": 1, "attn.to_v": 2}
|
||||
delta = torch.cat([d for _, d in sorted(parts, key=lambda pd: order[pd[0]])], 0)
|
||||
else:
|
||||
(part, delta), = parts
|
||||
if part == "ff.net.0.proj":
|
||||
up, gate = delta.chunk(2, 0)
|
||||
delta = torch.cat([gate, up], 0)
|
||||
want = F.linear(inp, weights[name + ".weight"] + delta, weights.get(name + ".bias"))
|
||||
with controller.activate(hop, sigma), controller.forward_hooks(controller.current()):
|
||||
got = layer(inp)
|
||||
torch.testing.assert_close(got, want, rtol=1e-4, atol=1e-5)
|
||||
torch.testing.assert_close(layer(inp), F.linear(inp, weights[name + ".weight"], weights.get(name + ".bias")),
|
||||
rtol=1e-4, atol=1e-5) # nothing stays hooked after the hop
|
||||
|
||||
|
||||
def test_a_hop_must_run_at_its_trained_time(ptd, folder):
|
||||
base = _tiny.native_model()
|
||||
_tiny.epfl_file(folder / "a.safetensors", base.get_model_object("diffusion_model"))
|
||||
controller, _ = _install(ptd, _clone(base), "a.safetensors")
|
||||
with pytest.raises(ValueError, match="trained time"):
|
||||
with controller.activate(1, 0.5):
|
||||
pass
|
||||
with pytest.raises(RuntimeError, match="no PTD hop is running"):
|
||||
controller.current()
|
||||
|
||||
|
||||
def test_the_chosen_sampler_is_replaced_by_the_adapters_own_hops(ptd, folder):
|
||||
base = _tiny.native_model()
|
||||
native = base.get_model_object("diffusion_model")
|
||||
_tiny.epfl_file(folder / "a.safetensors", native)
|
||||
_tiny.epfl_file(folder / "zero.safetensors", native, zero_a=True)
|
||||
seen = []
|
||||
for name in ("a.safetensors", "zero.safetensors"):
|
||||
model = _clone(base)
|
||||
controller, note = _install(ptd, model, name)
|
||||
assert "4 hops" in note
|
||||
calls = []
|
||||
original = controller.activate
|
||||
|
||||
def spy(hop, sigma, original=original, calls=calls):
|
||||
calls.append(hop)
|
||||
return original(hop, sigma)
|
||||
|
||||
controller.activate = spy
|
||||
# twenty euler steps were asked for; PTD runs its four hops
|
||||
out = _sample(model)
|
||||
assert calls == [0, 1, 2, 3]
|
||||
assert all(torch.isfinite(t).all() for t in out)
|
||||
seen.append(out)
|
||||
assert not torch.allclose(seen[0][0], seen[1][0])
|
||||
# the base model is not left patched
|
||||
plain = _sample(_clone(base))
|
||||
assert not torch.allclose(plain[0], seen[0][0])
|
||||
|
||||
|
||||
@pytest.mark.parametrize("curve", [False, True], ids=["full", "pruned"])
|
||||
def test_exact_and_fast_blocks_agree(ptd, folder, curve):
|
||||
base = _tiny.native_model(curve)
|
||||
_tiny.pdd_file(folder / "r.safetensors", base.get_model_object("diffusion_model"))
|
||||
exact = _sample(_installed(ptd, base, "exact"))
|
||||
fast = _sample(_installed(ptd, base, "fast"))
|
||||
for e, f in zip(exact, fast):
|
||||
assert torch.isfinite(e).all()
|
||||
assert torch.allclose(e, f, rtol=1e-4, atol=1e-4)
|
||||
|
||||
|
||||
def _installed(ptd, base, kernels):
|
||||
model = _clone(base)
|
||||
_install(ptd, model, "r.safetensors", kernels)
|
||||
return model
|
||||
|
||||
|
||||
def test_int8_ref2va_stays_close_to_float(ptd, folder):
|
||||
base = _tiny.native_model()
|
||||
native = base.get_model_object("diffusion_model")
|
||||
_tiny.pdd_file(folder / "r.safetensors", native)
|
||||
_tiny.pdd_file(folder / "q.safetensors", native, int8=True)
|
||||
ptd._kept.clear()
|
||||
full = _sample(_installed_file(ptd, base, "r.safetensors"))
|
||||
small = _sample(_installed_file(ptd, base, "q.safetensors"))
|
||||
for f, q in zip(full, small):
|
||||
assert torch.allclose(f, q, rtol=1e-2, atol=1e-2)
|
||||
|
||||
|
||||
def _installed_file(ptd, base, name):
|
||||
model = _clone(base)
|
||||
_install(ptd, model, name)
|
||||
return model
|
||||
|
||||
|
||||
def test_ref2va_heads_replace_the_native_heads_for_a_hop_only(ptd, folder):
|
||||
base = _tiny.native_model()
|
||||
native = base.get_model_object("diffusion_model")
|
||||
_tiny.pdd_file(folder / "r.safetensors", native)
|
||||
controller, _ = _install(ptd, _clone(base), "r.safetensors")
|
||||
head = native.final_layer.video_out
|
||||
x = torch.randn(5, head.in_features)
|
||||
plain = head(x)
|
||||
hop = 1
|
||||
with controller.activate(hop, float(controller.grid[hop])), controller.forward_hooks(controller.current()):
|
||||
hopped = head(x)
|
||||
weight, bias = controller.adapter.heads[hop]["final_layer.video_out"]
|
||||
torch.testing.assert_close(hopped, F.linear(x, weight, bias), atol=1e-5, rtol=1e-5)
|
||||
assert not torch.allclose(hopped, plain, atol=1e-3)
|
||||
assert torch.equal(head(x), plain)
|
||||
|
||||
|
||||
def test_grids_are_the_published_ones(ptd, folder):
|
||||
from modules.sampling.ptd import schedule
|
||||
assert schedule.RECIPES["h3_4step"][:4] == HOP_TIMES
|
||||
for hops in (4, 8):
|
||||
grid = schedule.pdd_grid(hops)
|
||||
assert len(grid) == hops + 1 and abs(grid[0] - schedule.T_MAX) < 1e-12 and grid[-1] == 0.0
|
||||
assert all(a > b for a, b in zip(grid, grid[1:]))
|
||||
for video, audio in schedule.pdd_plans(hops):
|
||||
for plan in (video, audio):
|
||||
assert abs(float(plan.sum()) - 1) < 1e-6 and int((plan > 0).sum()) == 32 // hops
|
||||
|
||||
|
||||
def test_a_file_for_another_model_is_refused_in_plain_words(ptd, folder):
|
||||
big = _tiny.native_model()
|
||||
_tiny.epfl_file(folder / "a.safetensors", big.get_model_object("diffusion_model"))
|
||||
smaller = _tiny.native_model(num_layers=1)
|
||||
with pytest.raises(ValueError, match="does not match the loaded model's layers"):
|
||||
_install(ptd, _clone(smaller), "a.safetensors")
|
||||
|
||||
|
||||
def test_not_a_ptd_file_is_refused(ptd, folder):
|
||||
from safetensors.torch import save_file
|
||||
save_file({"x": torch.zeros(2)}, str(folder / "n.safetensors"), metadata={"format": "pt"})
|
||||
with pytest.raises(ValueError, match="Not a PTD adapter file"):
|
||||
_install(ptd, _clone(_tiny.native_model()), "n.safetensors")
|
||||
|
||||
|
||||
def test_no_file_chosen_is_said_not_quarantined(ptd, folder):
|
||||
from modules._core import control
|
||||
with pytest.raises(control.Unavailable):
|
||||
ptd.install(_guarded(_clone(_tiny.native_model())), {"enabled": True, "adapter": "", "kernels": "exact"},
|
||||
"funpack.ptd")
|
||||
|
||||
|
||||
def test_off_installs_nothing(ptd, folder):
|
||||
patcher = _clone(_tiny.native_model())
|
||||
assert ptd.install(_guarded(patcher), {"enabled": False, "adapter": "", "kernels": "exact"}, "funpack.ptd") is None
|
||||
assert "funpack_ptd_controller" not in patcher.model_options
|
||||
|
||||
|
||||
def test_load_modifiers_strips_ptd_so_a_second_pass_reinstalls_it_once(ptd, folder):
|
||||
from modules._core import patching
|
||||
base = _tiny.native_model()
|
||||
_tiny.epfl_file(folder / "a.safetensors", base.get_model_object("diffusion_model"))
|
||||
model = _clone(base)
|
||||
_install(ptd, model, "a.safetensors")
|
||||
again = model.clone()
|
||||
assert patching.strip(again, "funpack.") >= 3
|
||||
assert "funpack_ptd_controller" not in again.model_options
|
||||
assert again.model_options.get("model_function_wrapper") is None
|
||||
_install(ptd, again, "a.safetensors") # no "already carries" error
|
||||
assert "funpack_ptd_controller" in again.model_options
|
||||
assert "funpack_ptd_controller" not in base.model_options and not base.object_patches
|
||||
|
||||
|
||||
def test_a_latent_that_does_not_start_empty_is_refused(ptd, folder):
|
||||
from modules.sampling.ptd.sampler import _refuse_unsupported
|
||||
|
||||
class Model:
|
||||
latent_image = torch.ones(1, 2)
|
||||
|
||||
with pytest.raises(ValueError, match="empty latent"):
|
||||
_refuse_unsupported(Model(), {})
|
||||
with pytest.raises(ValueError, match="masked"):
|
||||
_refuse_unsupported(object(), {"denoise_mask": torch.ones(1)})
|
||||
|
||||
|
||||
def test_it_installs_through_load_modifiers_on_an_h3_model(ptd, folder, monkeypatch):
|
||||
from core import registry
|
||||
from modules.sampling.modifiers.nodes import FunPackLoadModifiers
|
||||
base = _tiny.native_model()
|
||||
_tiny.pdd_file(folder / "r.safetensors", base.get_model_object("diffusion_model"))
|
||||
spec = registry.current().specs["ptd"]
|
||||
monkeypatch.setitem(spec.settings["adapter"], "options", [{"value": "r.safetensors", "label": "r"}])
|
||||
monkeypatch.setitem(spec.settings["adapter"], "default", "r.safetensors")
|
||||
patched, status = FunPackLoadModifiers.execute(_clone(base), {"ptd": {"enabled": True}}).result
|
||||
assert "ptd: r.safetensors, 4 hops" in status
|
||||
assert "failed" not in status
|
||||
out = _sample(patched)
|
||||
assert all(torch.isfinite(t).all() for t in out)
|
||||
# switched off, the same model samples as the plain one does
|
||||
plain, status = FunPackLoadModifiers.execute(_clone(base), {"ptd": {"enabled": False}}).result
|
||||
assert "ptd" not in status.splitlines()[0]
|
||||
|
||||
|
||||
def test_an_ordinary_lora_stacks_with_ptd_in_either_order(ptd, folder):
|
||||
import comfy.sd
|
||||
base = _tiny.native_model()
|
||||
native = base.get_model_object("diffusion_model")
|
||||
dense = _tiny.epfl_file(folder / "a.safetensors", native)
|
||||
g = torch.Generator().manual_seed(2)
|
||||
layer = native.get_submodule("blocks.1.mlp.fc2")
|
||||
up, down = torch.randn(layer.out_features, 4, generator=g) * 0.1, torch.randn(4, layer.in_features, generator=g) * 0.1
|
||||
lora = {"blocks.1.mlp.fc2.lora_up.weight": up, "blocks.1.mlp.fc2.lora_down.weight": down}
|
||||
weight = layer.weight.detach().clone()
|
||||
for lora_first in (True, False):
|
||||
ptd._kept.clear()
|
||||
model = comfy.sd.load_lora_for_models(_clone(base), None, lora, 0.7, 0)[0] if lora_first else _clone(base)
|
||||
controller, _ = _install(ptd, model, "a.safetensors")
|
||||
if not lora_first:
|
||||
model = comfy.sd.load_lora_for_models(model, None, lora, 0.7, 0)[0]
|
||||
import comfy.model_management
|
||||
comfy.model_management.load_models_gpu([model])
|
||||
try:
|
||||
inp = torch.randn(5, layer.in_features)
|
||||
hop = 2
|
||||
want = F.linear(inp, weight + 0.7 * up @ down + dense[("blocks.1.mlp.fc2", "ff.net.2", hop)],
|
||||
layer.bias)
|
||||
with controller.activate(hop, HOP_TIMES[hop]), controller.forward_hooks(controller.current()):
|
||||
torch.testing.assert_close(layer(inp), want, rtol=1e-4, atol=1e-5)
|
||||
finally:
|
||||
comfy.model_management.unload_all_models()
|
||||
Reference in New Issue
Block a user