From 3f8a007da2393ea4c1939774397b905a06f0fb9d Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 23 Oct 2026 18:12:38 +0000 Subject: [PATCH] fix: PTD hop schedule keeps the sampler's device and dtype The schedule handed to step-aware code (SLA, H3's head bank) was a CPU copy while the call's sigma sits on the GPU, which fails with "Expected all tensors to be on the same device". It now matches the schedule the sampler was given. Test uses the meta device as a stand-in for CUDA. Co-Authored-By: Claude Sonnet 5.5 Claude-Session: https://claude.ai/code/session_0164nckPWA2L2ugVqiDcXzc3 --- modules/sampling/ptd/__init__.py | 7 +++++-- modules/sampling/ptd/tests/test_ptd.py | 27 ++++++++++++++++++++++++++ 2 files changed, 32 insertions(+), 2 deletions(-) diff --git a/modules/sampling/ptd/__init__.py b/modules/sampling/ptd/__init__.py index e2149f6c..be00b316 100644 --- a/modules/sampling/ptd/__init__.py +++ b/modules/sampling/ptd/__init__.py @@ -102,13 +102,16 @@ def _swap(controller): except Exception: # noqa: BLE001 -- only the progress bar and previews are lost pass # Step-aware modules (SLA, momentum, velocity bias) read the schedule from here: give them the real one. + # On the device and dtype the sampler was given: ComfyUI's own schedule sits on the GPU, and a model that + # compares it with its call's sigma (H3's head bank, SLA) fails on a CPU copy. + grid = controller.grid.to(device=sigmas.device, dtype=sigmas.dtype) options = extra_args.get("model_options", {}).get("transformer_options") if options is not None: - options["sample_sigmas"] = controller.grid.float() + options["sample_sigmas"] = grid 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, + return ptd.sample(guider, grid, extra_args, callback, noise, latent_image, denoise_mask, disable_pbar) return run diff --git a/modules/sampling/ptd/tests/test_ptd.py b/modules/sampling/ptd/tests/test_ptd.py index aadce67e..4fcb3bf8 100644 --- a/modules/sampling/ptd/tests/test_ptd.py +++ b/modules/sampling/ptd/tests/test_ptd.py @@ -330,3 +330,30 @@ def test_the_file_list_is_read_again_when_settings_open(ptd, folder, monkeypatch 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]["funpack.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