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