diff --git a/modules/sampling/two_stage/stages.py b/modules/sampling/two_stage/stages.py index 0cf6e88c..ef4086ac 100644 --- a/modules/sampling/two_stage/stages.py +++ b/modules/sampling/two_stage/stages.py @@ -45,7 +45,8 @@ def keep_whole_schedule(executor, guider, sigmas, extra_args, *args, **kwargs): ramp, step windows) sees 8 steps and not two runs of 4.""" full = getattr(guider, FULL, None) if full is not None: - extra_args["model_options"].setdefault("transformer_options", {})["sample_sigmas"] = full + extra_args["model_options"].setdefault("transformer_options", {})["sample_sigmas"] = \ + full.to(device=sigmas.device, dtype=sigmas.dtype) # where this stage's own schedule lives (the GPU) return executor(guider, sigmas, extra_args, *args, **kwargs) diff --git a/modules/sampling/two_stage/tests/test_two_stage.py b/modules/sampling/two_stage/tests/test_two_stage.py index 73862866..56e6a940 100644 --- a/modules/sampling/two_stage/tests/test_two_stage.py +++ b/modules/sampling/two_stage/tests/test_two_stage.py @@ -76,3 +76,17 @@ def test_every_stage_sees_the_whole_schedule(model): sigmas = (1.0, 0.9, 0.7, 0.3, 0.0) _run(model, {"enabled": True, "size": 0.5, "steps": 2}, sigmas=sigmas, seen=seen) assert len(seen) == 4 and all(s == pytest.approx(list(sigmas)) for s in seen) + + +def test_the_whole_schedule_lives_where_the_stage_schedule_does(): + from ..stages import FULL, keep_whole_schedule + + class Guider: + pass + + guider, extra = Guider(), {"model_options": {}} + setattr(guider, FULL, torch.tensor([1.0, 0.5, 0.0], dtype=torch.float64)) + stage = torch.zeros(2, device="meta") + keep_whole_schedule(lambda *a, **k: None, guider, stage, extra) + got = extra["model_options"]["transformer_options"]["sample_sigmas"] + assert got.device == stage.device and got.dtype == stage.dtype