Claude
3adebb731f
Old project/workflow names and browser settings migrate on load; /funpack redirects to /fun. Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_016acRss69eH1235eRdFTZHY
360 lines
16 KiB
Python
360 lines
16 KiB
Python
"""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, "funui.ptd", patching.Dropped())
|
|
|
|
|
|
def _install(ptd, patcher, name, kernels="exact"):
|
|
note = ptd.install(_guarded(patcher), {"enabled": True, "adapter": name, "kernels": kernels}, "funui.ptd")
|
|
return patcher.model_options["funui_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"},
|
|
"funui.ptd")
|
|
|
|
|
|
def test_off_installs_nothing(ptd, folder):
|
|
patcher = _clone(_tiny.native_model())
|
|
assert ptd.install(_guarded(patcher), {"enabled": False, "adapter": "", "kernels": "exact"}, "funui.ptd") is None
|
|
assert "funui_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, "funui.") >= 3
|
|
assert "funui_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 "funui_ptd_controller" in again.model_options
|
|
assert "funui_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 FunUILoadModifiers
|
|
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 = FunUILoadModifiers.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 = FunUILoadModifiers.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()
|
|
|
|
|
|
def test_step_aware_modules_are_given_the_hop_schedule(ptd, folder):
|
|
base = _tiny.native_model()
|
|
_tiny.epfl_file(folder / "a.safetensors", base.get_model_object("diffusion_model"))
|
|
model = _clone(base)
|
|
seen = []
|
|
|
|
def probe(apply_fn, args):
|
|
seen.append(args["c"]["transformer_options"]["sample_sigmas"].tolist())
|
|
return apply_fn(args["input"], args["timestep"], **args["c"])
|
|
|
|
model.set_model_unet_function_wrapper(probe)
|
|
_install(ptd, model, "a.safetensors")
|
|
_sample(model) # twenty steps asked for
|
|
assert seen and all(len(s) == 5 for s in seen)
|
|
|
|
|
|
def test_the_file_list_is_read_again_when_settings_open(ptd, folder, monkeypatch):
|
|
pytest.importorskip("aiohttp")
|
|
from core import registry, routes
|
|
spec = registry.current().specs["ptd"]
|
|
monkeypatch.setitem(spec.settings["adapter"], "options", [{"value": "", "label": "No file in models/ptd"}])
|
|
monkeypatch.setitem(spec.settings["adapter"], "default", "")
|
|
base = _tiny.native_model()
|
|
_tiny.epfl_file(folder / "new.safetensors", base.get_model_object("diffusion_model"))
|
|
listed = {m["id"]: m for m in routes.manifest()["modules"]}
|
|
adapter = listed["ptd"]["settings"]["adapter"]
|
|
assert "new.safetensors" in [o["value"] for o in adapter["options"]]
|
|
assert adapter["default"] in [o["value"] for o in adapter["options"]]
|
|
|
|
|
|
def test_the_hop_schedule_sits_where_the_sampler_was_given_its_own(ptd, folder, monkeypatch):
|
|
"""On a GPU ComfyUI hands the sampler a CUDA schedule; the one PTD swaps in must be on the same device
|
|
(a CPU copy met the call's CUDA sigma: "Expected all tensors to be on the same device"). The meta device
|
|
stands in for CUDA here."""
|
|
import comfy.samplers
|
|
from comfy.patcher_extension import WrappersMP
|
|
base = _tiny.native_model()
|
|
_tiny.epfl_file(folder / "a.safetensors", base.get_model_object("diffusion_model"))
|
|
model = _clone(base)
|
|
_install(ptd, model, "a.safetensors")
|
|
given = {}
|
|
monkeypatch.setattr(comfy.samplers.KSAMPLER, "sample",
|
|
lambda self, wrap, sigmas, extra_args, *a, **k: given.update(sigmas=sigmas) or "done")
|
|
|
|
class Guider:
|
|
cfg = 1.0
|
|
model_patcher = None
|
|
|
|
extra = {"model_options": {"transformer_options": {}}}
|
|
run = model.wrappers[WrappersMP.SAMPLER_SAMPLE]["funui.ptd"][0]
|
|
sigmas = torch.linspace(1.0, 0.0, 21, device="meta")
|
|
assert run(None, Guider(), sigmas, extra, None, None, None, None, True) == "done"
|
|
assert given["sigmas"].device == sigmas.device and given["sigmas"].dtype == sigmas.dtype
|
|
seen = extra["model_options"]["transformer_options"]["sample_sigmas"]
|
|
assert seen.device == sigmas.device and len(seen) == 5
|