diff --git a/CHANGELOG.md b/CHANGELOG.md index e1667a44..cef200f6 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -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. diff --git a/README.md b/README.md index 48790a7e..a18dd3e2 100644 --- a/README.md +++ b/README.md @@ -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. diff --git a/modules/sampling/modifiers/tests/test_every_setting.py b/modules/sampling/modifiers/tests/test_every_setting.py index 407e0ecb..06d16178 100644 --- a/modules/sampling/modifiers/tests/test_every_setting.py +++ b/modules/sampling/modifiers/tests/test_every_setting.py @@ -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): diff --git a/modules/sampling/ptd/LICENSE-Apache-2.0.txt b/modules/sampling/ptd/LICENSE-Apache-2.0.txt new file mode 100644 index 00000000..d6456956 --- /dev/null +++ b/modules/sampling/ptd/LICENSE-Apache-2.0.txt @@ -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. diff --git a/modules/sampling/ptd/NOTICE b/modules/sampling/ptd/NOTICE new file mode 100644 index 00000000..4099e704 --- /dev/null +++ b/modules/sampling/ptd/NOTICE @@ -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. diff --git a/modules/sampling/ptd/__init__.py b/modules/sampling/ptd/__init__.py new file mode 100644 index 00000000..6d5787d2 --- /dev/null +++ b/modules/sampling/ptd/__init__.py @@ -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} diff --git a/modules/sampling/ptd/adapters.py b/modules/sampling/ptd/adapters.py new file mode 100644 index 00000000..1afaa646 --- /dev/null +++ b/modules/sampling/ptd/adapters.py @@ -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..lora_A..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.+)\.lora_(?P[AB])\.(?P0|[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) diff --git a/modules/sampling/ptd/controller.py b/modules/sampling/ptd/controller.py new file mode 100644 index 00000000..7b3928a5 --- /dev/null +++ b/modules/sampling/ptd/controller.py @@ -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" diff --git a/modules/sampling/ptd/sampler.py b/modules/sampling/ptd/sampler.py new file mode 100644 index 00000000..4d715928 --- /dev/null +++ b/modules/sampling/ptd/sampler.py @@ -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 diff --git a/modules/sampling/ptd/schedule.py b/modules/sampling/ptd/schedule.py new file mode 100644 index 00000000..270766c0 --- /dev/null +++ b/modules/sampling/ptd/schedule.py @@ -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]) diff --git a/modules/sampling/ptd/tests/_tiny.py b/modules/sampling/ptd/tests/_tiny.py new file mode 100644 index 00000000..38cd9764 --- /dev/null +++ b/modules/sampling/ptd/tests/_tiny.py @@ -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) diff --git a/modules/sampling/ptd/tests/test_ptd.py b/modules/sampling/ptd/tests/test_ptd.py new file mode 100644 index 00000000..1345d5be --- /dev/null +++ b/modules/sampling/ptd/tests/test_ptd.py @@ -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()