From e8d8fcc81d5baa82a0df65cd2342a264892db4dd Mon Sep 17 00:00:00 2001 From: Boobkake22 Date: Sat, 27 Jun 2026 15:09:06 -0700 Subject: [PATCH] Add progressive handoff sampling mode --- README.md | 74 ++++++++++++++++-- nodes.py | 146 +++++++++++++++++++++++++++++++++++- tests/test_nodes_smoke.py | 69 +++++++++++++++++ tests/test_wan22_planner.py | 29 +++++++ wan22_planner.py | 128 ++++++++++++++++++++++++++++++- 5 files changed, 436 insertions(+), 10 deletions(-) diff --git a/README.md b/README.md index 39eafd6..31590ec 100644 --- a/README.md +++ b/README.md @@ -56,6 +56,22 @@ Outputs a priority combo compatible with Sampling Plan (Wan 2.2): - `Motion / Structure` - `Detail / Refinement` +### Handoff Selector + +Outputs a handoff-mode combo compatible with Sampling Plan (Wan 2.2): + +- `Direct Latent` +- `Progressive Transcode` + +`Direct Latent` is the normal two-stage continuation: the high sampler stops at +the handoff sigma, the low sampler starts from that latent, and low-stage noise +is disabled. + +`Progressive Transcode` is for workflows that decode/upscale/re-encode between +the high and low stages. The high sampler receives a terminal high schedule for a +clean decoded image, and the low sampler re-noises the re-encoded latent at the +handoff sigma. + ### Step Budget Defines two full-range-equivalent budgets: @@ -134,6 +150,7 @@ Managed controls: - **Step Budget:** accelerated and full-range-equivalent steps - **Scheduler:** ComfyUI sigma scheduler - **Priority:** Balanced, 50/50, Motion / Structure, or Detail / Refinement +- **Handoff Mode:** Direct Latent or Progressive Transcode The connected `MODEL` is used to calculate the real ComfyUI sigma schedule. It can be either Wan expert model. @@ -146,9 +163,11 @@ The planner uses the official Wan 2.2 boundaries: - T2V: `0.875` - I2V: `0.900` -Priority determines the high/low share of the denoising range. Acceleration -determines which budget is projected onto each share. The resulting high and low -step counts are added to produce the effective total schedule. +Priority determines the high/low share of the denoising range. Acceleration and +Handoff Mode determine which budget is projected onto each share. Direct Latent +projects active accelerated/full budgets independently. Progressive Transcode +uses one strategy budget across both stages: the full budget when acceleration is +off, or the accelerated budget when any acceleration is active. `Auto` selects an evidence-backed curve profile: @@ -210,6 +229,39 @@ Sampling Plan ────────────┘ Model Pair Breakout has no user-facing widgets. Shift is owned by the plan, including any Shift Override. +### Custom Sampler Breakout + +Connect the plan and a seed. The node outputs: + +- high noise +- low noise +- high sigmas +- low sigmas +- total/high/low step counts +- shift +- the plan passthrough + +Use this with `SamplerCustomAdvanced` and `CFGGuider`. + +In `Direct Latent`, high noise is random, low noise is disabled, and the high +sigmas end at the same handoff sigma where the low sigmas begin. + +In `Progressive Transcode`, high noise is random, low noise is random, and the +high sigmas terminate at zero for a clean high-stage decode. The low sigmas still +begin at the original handoff sigma, so the re-encoded latent is re-noised at the +same point where low-noise refinement begins. + +Recommended progressive wiring: + +```text +Custom Sampler Breakout noise_high → high SamplerCustomAdvanced noise +Custom Sampler Breakout sigmas_high → high SamplerCustomAdvanced sigmas +high SamplerCustomAdvanced output → decode/upscale/re-encode +Custom Sampler Breakout noise_low → low SamplerCustomAdvanced noise +Custom Sampler Breakout sigmas_low → low SamplerCustomAdvanced sigmas +re-encoded latent → low SamplerCustomAdvanced latent_image +``` + ### KSampler Breakout Connect the plan, high/low expert models, and the sampler selected for the @@ -241,6 +293,11 @@ using piecewise-resampled sigmas are not KSampler-representable; the node report that incompatibility and directs the workflow to Sigma Breakout instead of silently approximating the curve. +Progressive Transcode plans also bypass KSampler Breakout because that node does +not emit the required add-noise and terminal-output controls. Use Custom Sampler +Breakout with SamplerCustomAdvanced, or wire KSampler Advanced's add-noise and +return-with-leftover-noise widgets manually. + For parity testing against older Wan 2.2 workflows, prefer this KSampler Breakout path whenever the plan reports `curve_mode: exact`. It lets ComfyUI's native KSampler Advanced implementation own `start_at_step`, `end_at_step`, and @@ -257,7 +314,10 @@ Exposes: - total/high/low step counts - shift -The high schedule ends at the same sigma where the low schedule begins. +In Direct Latent, the high schedule ends at the same sigma where the low schedule +begins. In Progressive Transcode, Sigma Breakout emits the terminal high schedule +for a clean decode/upscale/re-encode handoff, while the low schedule still begins +at the original handoff sigma. Sigma Breakout is the authoritative sigma-output path for every plan. It preserves the exact validated curve, including mixed-budget and @@ -265,7 +325,7 @@ priority-adjusted piecewise curves that ordinary KSampler controls cannot express. Use it when KSampler Breakout reports that the plan cannot be represented by native KSampler Advanced controls. -For `SamplerCustomAdvanced`, continue the two expert stages as follows: +For direct `SamplerCustomAdvanced`, continue the two expert stages as follows: ```text RandomNoise → high SamplerCustomAdvanced → output → low SamplerCustomAdvanced → output @@ -281,7 +341,9 @@ For KSampler parity, use `CFGGuider` for each custom-sampler stage, feed the same positive/negative conditioning and CFG values that the KSampler Advanced nodes used, use the same `KSamplerSelect` sampler, and do not insert a separate `BasicScheduler`/scheduler node between the plan and the samplers. Sigma -Breakout already emits the high and low sigma slices. +Breakout already emits the mode-correct high and low sigma slices. Custom +Sampler Breakout adds the matching noise objects so the workflow needs less +manual switching. See [Comfy sampler parity notes](docs/comfy-sampler-parity.md) for the source comparison between KSampler Advanced and SamplerCustomAdvanced. diff --git a/nodes.py b/nodes.py index 34c1b4e..93d2dcd 100644 --- a/nodes.py +++ b/nodes.py @@ -4,10 +4,14 @@ from typing import Any import comfy.samplers from comfy_extras.nodes_model_advanced import ModelSamplingSD3 +from comfy_extras.nodes_custom_sampler import Noise_EmptyNoise, Noise_RandomNoise from .wan22_planner import ( ACCELERATION_BOTH, ACCELERATION_OPTIONS, + HANDOFF_DIRECT, + HANDOFF_OPTIONS, + HANDOFF_PROGRESSIVE, PLAN_TYPE, PRIORITY_BALANCED, PRIORITY_OPTIONS, @@ -91,6 +95,12 @@ class _PrioritySelectorReturnTypes: return (PRIORITY_OPTIONS,) +class _HandoffSelectorReturnTypes: + def __get__(self, instance, owner): + del instance, owner + return (HANDOFF_OPTIONS,) + + class _AccelerationSelectorReturnTypes: def __get__(self, instance, owner): del instance, owner @@ -124,6 +134,7 @@ def _acceleration_state(plan: dict[str, Any]) -> dict[str, Any]: return { "type": ACCELERATION_STATE_TYPE, "mode": plan["acceleration"], + "handoff_mode": plan["handoff_mode"], "accelerate_high": plan["accelerate_high"], "accelerate_low": plan["accelerate_low"], } @@ -141,6 +152,7 @@ def _rebuild_plan(plan: Any, **override_updates: Any) -> dict[str, Any]: full_steps=source["full_steps"], scheduler=source["scheduler"], priority=source["priority"], + handoff_mode=source.get("handoff_mode", HANDOFF_DIRECT), sigma_provider=_sigma_provider(source["model"]), forced_steps_high=overrides.get("steps_high"), forced_shift=overrides.get("shift"), @@ -226,6 +238,18 @@ class SamplingPlanWan22: ), }, ), + "handoff_mode": ( + HANDOFF_OPTIONS, + { + "default": HANDOFF_DIRECT, + "tooltip": ( + "Direct Latent continues the high latent into the low " + "stage with low noise disabled. Progressive Transcode " + "cleans the high output for decode/upscale/re-encode, " + "then re-noises the low stage at the handoff." + ), + }, + ), } } @@ -242,6 +266,7 @@ class SamplingPlanWan22: step_budget, scheduler, priority, + handoff_mode=HANDOFF_DIRECT, ): step_budget = _validate_step_budget(step_budget) plan = build_plan( @@ -252,6 +277,7 @@ class SamplingPlanWan22: full_steps=step_budget["full_steps"], scheduler=scheduler, priority=priority, + handoff_mode=handoff_mode, sigma_provider=_sigma_provider(model), ) print(f"Sampling Plan (Wan 2.2): {plan['summary']}") @@ -399,6 +425,40 @@ class PrioritySelector: return (priority,) +class HandoffSelector: + DESCRIPTION = ( + "Selects how the high-noise stage hands off to the low-noise stage." + ) + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "handoff_mode": ( + HANDOFF_OPTIONS, + { + "default": HANDOFF_DIRECT, + "tooltip": ( + "Direct Latent is a normal high-to-low continuation. " + "Progressive Transcode is for decode/upscale/re-encode " + "between stages and re-noises the low stage." + ), + }, + ) + } + } + + RETURN_TYPES = _HandoffSelectorReturnTypes() + RETURN_NAMES = ("handoff_mode",) + FUNCTION = "select_handoff" + CATEGORY = HELPER_CATEGORY + + def select_handoff(self, handoff_mode): + if handoff_mode not in HANDOFF_OPTIONS: + raise ValueError(f"Unknown handoff mode: {handoff_mode}") + return (handoff_mode,) + + class StepBudget: DESCRIPTION = ( "Defines full-range-equivalent step budgets for accelerated and " @@ -761,6 +821,79 @@ class ModelPairBreakout: ) +class CustomSamplerBreakout: + DESCRIPTION = ( + "Exposes the noise objects and sigma schedules for a two-stage " + "SamplerCustomAdvanced Wan 2.2 branch." + ) + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "plan": (PLAN_TYPE,), + "noise_seed": ( + "INT", + { + "default": 0, + "min": 0, + "max": 0xFFFFFFFFFFFFFFFF, + "control_after_generate": True, + "tooltip": ( + "Seed used for the high stage, and for the low stage " + "when Progressive Transcode needs to re-noise." + ), + }, + ), + } + } + + RETURN_TYPES = ( + "NOISE", + "NOISE", + "SIGMAS", + "SIGMAS", + "INT", + "INT", + "INT", + "FLOAT", + PLAN_TYPE, + ) + RETURN_NAMES = ( + "noise_high", + "noise_low", + "sigmas_high", + "sigmas_low", + "steps", + "steps_high", + "steps_low", + "shift", + "plan", + ) + FUNCTION = "breakout" + CATEGORY = CATEGORY + + def breakout(self, plan, noise_seed): + plan = validate_plan(plan) + if plan["noise_low_mode"] == "random": + noise_low = Noise_RandomNoise(int(noise_seed)) + else: + noise_low = Noise_EmptyNoise() + + return _ui_result( + plan, + Noise_RandomNoise(int(noise_seed)), + noise_low, + plan["sigmas_high_sampling"], + plan["sigmas_low"], + plan["steps"], + plan["steps_high"], + plan["steps_low"], + float(plan["shift"]), + plan, + ) + + class KSamplerBreakout: DESCRIPTION = ( "Prepares a high/low Wan 2.2 model pair and exposes the coordinated " @@ -820,6 +953,13 @@ class KSamplerBreakout: def breakout(self, plan, model_high, model_low, sampler): plan = validate_plan(plan) + if plan.get("handoff_mode") == HANDOFF_PROGRESSIVE: + raise ValueError( + "Progressive Transcode needs terminal high output and low-stage " + "re-noise controls. Use Custom Sampler Breakout with " + "SamplerCustomAdvanced, or wire KSampler Advanced add_noise / " + "return_with_leftover_noise manually." + ) if plan.get("curve_mode") != "exact": raise ValueError( "This plan uses a piecewise sigma curve and cannot be represented " @@ -891,7 +1031,7 @@ class SigmaBreakout: plan = validate_plan(plan) return _ui_result( plan, - plan["sigmas_high"], + plan["sigmas_high_sampling"], plan["sigmas_low"], plan["sigmas"], plan["steps"], @@ -907,6 +1047,7 @@ NODE_CLASS_MAPPINGS = { "SamplerSelector": SamplerSelector, "TaskSelector": TaskSelector, "PrioritySelector": PrioritySelector, + "HandoffSelector": HandoffSelector, "StepBudget": StepBudget, "AccelerationSelector": AccelerationSelector, "AccelerationModelPair": AccelerationModelPair, @@ -915,6 +1056,7 @@ NODE_CLASS_MAPPINGS = { "RangeSplitOverrideLegacy": RangeSplitOverrideLegacy, "ShiftOverride": ShiftOverride, "ModelPairBreakout": ModelPairBreakout, + "CustomSamplerBreakout": CustomSamplerBreakout, "KSamplerBreakout": KSamplerBreakout, "SigmaBreakout": SigmaBreakout, } @@ -924,6 +1066,7 @@ NODE_DISPLAY_NAME_MAPPINGS = { "SamplerSelector": "Sampler Selector", "TaskSelector": "Task Selector", "PrioritySelector": "Priority Selector", + "HandoffSelector": "Handoff Selector", "StepBudget": "Step Budget", "AccelerationSelector": "Acceleration Selector", "AccelerationModelPair": "Acceleration Model Pair", @@ -932,6 +1075,7 @@ NODE_DISPLAY_NAME_MAPPINGS = { "RangeSplitOverrideLegacy": "Range Split Override (Legacy)", "ShiftOverride": "Shift Override", "ModelPairBreakout": "Model Pair Breakout", + "CustomSamplerBreakout": "Custom Sampler Breakout", "KSamplerBreakout": "KSampler Breakout", "SigmaBreakout": "Sigma Breakout", } diff --git a/tests/test_nodes_smoke.py b/tests/test_nodes_smoke.py index 5abd9f8..5687556 100644 --- a/tests/test_nodes_smoke.py +++ b/tests/test_nodes_smoke.py @@ -30,6 +30,15 @@ class FakeModelSamplingSD3: return (patched,) +class FakeNoiseRandomNoise: + def __init__(self, seed): + self.seed = int(seed) + + +class FakeNoiseEmptyNoise: + pass + + def fake_calculate_sigmas(model_sampling, scheduler, steps): del scheduler shift = model_sampling.shift @@ -76,11 +85,15 @@ def load_nodes_module(): comfy_extras.__path__ = [] model_advanced = types.ModuleType("comfy_extras.nodes_model_advanced") model_advanced.ModelSamplingSD3 = FakeModelSamplingSD3 + nodes_custom_sampler = types.ModuleType("comfy_extras.nodes_custom_sampler") + nodes_custom_sampler.Noise_RandomNoise = FakeNoiseRandomNoise + nodes_custom_sampler.Noise_EmptyNoise = FakeNoiseEmptyNoise sys.modules["comfy"] = comfy sys.modules["comfy.samplers"] = comfy_samplers sys.modules["comfy_extras"] = comfy_extras sys.modules["comfy_extras.nodes_model_advanced"] = model_advanced + sys.modules["comfy_extras.nodes_custom_sampler"] = nodes_custom_sampler spec = importlib.util.spec_from_file_location( PACKAGE_NAME, @@ -107,6 +120,7 @@ class NodeSmokeTests(unittest.TestCase): "SamplerSelector", "TaskSelector", "PrioritySelector", + "HandoffSelector", "StepBudget", "AccelerationSelector", "AccelerationModelPair", @@ -115,6 +129,7 @@ class NodeSmokeTests(unittest.TestCase): "RangeSplitOverrideLegacy", "ShiftOverride", "ModelPairBreakout", + "CustomSamplerBreakout", "KSamplerBreakout", "SigmaBreakout", }, @@ -177,6 +192,19 @@ class NodeSmokeTests(unittest.TestCase): ("50/50 Split",), ) + def test_handoff_selector_matches_plan_combo(self): + selector = self.nodes.HandoffSelector + selector_options = selector.INPUT_TYPES()["required"]["handoff_mode"][0] + plan_options = self.nodes.SamplingPlanWan22.INPUT_TYPES()["required"][ + "handoff_mode" + ][0] + self.assertIs(selector_options, plan_options) + self.assertEqual(selector.RETURN_TYPES, (plan_options,)) + self.assertEqual( + selector().select_handoff("Progressive Transcode"), + ("Progressive Transcode",), + ) + def test_acceleration_selector_matches_plan_combo(self): selector = self.nodes.AccelerationSelector options = selector.INPUT_TYPES()["required"]["acceleration"][0] @@ -377,6 +405,16 @@ class NodeSmokeTests(unittest.TestCase): revised["steps_low"] + 1, ) + custom_result = self.nodes.CustomSamplerBreakout().breakout( + revised, + 1234, + )["result"] + self.assertEqual(len(custom_result), 9) + self.assertIsInstance(custom_result[0], FakeNoiseRandomNoise) + self.assertIsInstance(custom_result[1], FakeNoiseEmptyNoise) + self.assertEqual(custom_result[0].seed, 1234) + self.assertEqual(list(custom_result[2]), list(revised["sigmas_high"])) + model_pair_result = self.nodes.ModelPairBreakout().breakout( revised, FakeModel(), @@ -421,6 +459,37 @@ class NodeSmokeTests(unittest.TestCase): "custom_sampler", ) + def test_progressive_custom_sampler_breakout_uses_terminal_high_and_renoise(self): + plan_node = self.nodes.SamplingPlanWan22() + budget = self.nodes.StepBudget().create_budget(10, 30)[0] + plan = plan_node.create_plan( + FakeModel(), + "I2V", + "Low only", + budget, + "simple", + "50/50 Split", + "Progressive Transcode", + )["result"][0] + self.assertEqual((plan["steps_high"], plan["steps_low"]), (5, 5)) + + result = self.nodes.CustomSamplerBreakout().breakout(plan, 9876)["result"] + self.assertIsInstance(result[0], FakeNoiseRandomNoise) + self.assertIsInstance(result[1], FakeNoiseRandomNoise) + self.assertEqual(result[0].seed, 9876) + self.assertEqual(result[1].seed, 9876) + self.assertEqual(result[2][-1], 0.0) + self.assertEqual(list(result[2]), list(plan["sigmas_high_terminal"])) + self.assertEqual(list(result[3]), list(plan["sigmas_low"])) + + with self.assertRaisesRegex(ValueError, "Progressive Transcode"): + self.nodes.KSamplerBreakout().breakout( + plan, + FakeModel(), + FakeModel(), + "euler", + ) + def test_override_nodes_expose_compact_exact_controls(self): split = self.nodes.StepSplitOverride.INPUT_TYPES()["required"] shift = self.nodes.ShiftOverride.INPUT_TYPES()["required"] diff --git a/tests/test_wan22_planner.py b/tests/test_wan22_planner.py index 20217c4..87eafa1 100644 --- a/tests/test_wan22_planner.py +++ b/tests/test_wan22_planner.py @@ -9,6 +9,8 @@ from wan22_planner import ( CURVE_PROFILE_COMFYUI, CURVE_PROFILE_DISTILLED, CURVE_PROFILE_NATIVE, + HANDOFF_DIRECT, + HANDOFF_PROGRESSIVE, PLAN_VERSION, PRIORITY_BALANCED, PRIORITY_DETAIL, @@ -47,6 +49,7 @@ def build(**updates): "full_steps": 30, "scheduler": "simple", "priority": PRIORITY_BALANCED, + "handoff_mode": HANDOFF_DIRECT, "sigma_provider": shifted_linear_provider, } arguments.update(updates) @@ -93,6 +96,32 @@ class ProfileSelectionTests(unittest.TestCase): self.assertEqual(list(plan["sigmas_low"]), sigmas[split:]) self.assertEqual(plan["sigmas_high"][-1], plan["sigmas_low"][0]) self.assertNotEqual(plan["sigmas_high"][-1], 0.0) + self.assertEqual(list(plan["sigmas_high_sampling"]), list(plan["sigmas_high"])) + self.assertEqual(plan["noise_low_mode"], "disable") + + def test_progressive_low_only_uses_strategy_budget_and_terminal_high(self): + plan = build( + acceleration=ACCELERATION_LOW, + priority=PRIORITY_EVEN, + handoff_mode=HANDOFF_PROGRESSIVE, + ) + self.assertEqual((plan["high_budget"], plan["low_budget"]), (10, 10)) + self.assertEqual((plan["steps_high"], plan["steps_low"]), (5, 5)) + self.assertEqual(plan["steps"], 10) + self.assertEqual(plan["handoff_mode"], HANDOFF_PROGRESSIVE) + self.assertEqual(plan["high_output_mode"], "terminal") + self.assertEqual(plan["noise_high_mode"], "random") + self.assertEqual(plan["noise_low_mode"], "random") + self.assertEqual(plan["sigmas_high"][-1], plan["sigmas_low"][0]) + self.assertNotEqual(plan["sigmas_high"][-1], 0.0) + self.assertEqual(plan["sigmas_high_terminal"][-1], 0.0) + self.assertEqual( + list(plan["sigmas_high_sampling"]), + list(plan["sigmas_high_terminal"]), + ) + self.assertTrue( + any("full-step budget is inactive" in item for item in plan["warnings"]) + ) def test_unaccelerated_uses_task_native_shift(self): t2v = build( diff --git a/wan22_planner.py b/wan22_planner.py index fc780fa..1354391 100644 --- a/wan22_planner.py +++ b/wan22_planner.py @@ -7,7 +7,7 @@ from typing import Any PLAN_TYPE = "WAN22_SAMPLING_PLAN" -PLAN_VERSION = 3 +PLAN_VERSION = 4 TASK_PROFILES = { "T2V": { @@ -43,6 +43,13 @@ PRIORITY_OPTIONS = [ PRIORITY_DETAIL, ] +HANDOFF_DIRECT = "Direct Latent" +HANDOFF_PROGRESSIVE = "Progressive Transcode" +HANDOFF_OPTIONS = [ + HANDOFF_DIRECT, + HANDOFF_PROGRESSIVE, +] + CURVE_PROFILE_NATIVE = "wan_native" CURVE_PROFILE_DISTILLED = "lightx2v_distilled" CURVE_PROFILE_COMFYUI = "comfyui_yaw" @@ -240,6 +247,33 @@ def project_stage_steps( return steps_high, steps_low, high_budget, low_budget +def project_handoff_steps( + *, + accelerated_steps: int, + full_steps: int, + acceleration: str, + high_range_share: float, + handoff_mode: str, +) -> tuple[int, int, int, int]: + if handoff_mode == HANDOFF_DIRECT: + return project_stage_steps( + accelerated_steps=accelerated_steps, + full_steps=full_steps, + acceleration=acceleration, + high_range_share=high_range_share, + ) + + if handoff_mode != HANDOFF_PROGRESSIVE: + raise ValueError(f"Unknown Wan 2.2 handoff mode: {handoff_mode}") + + strategy_budget = ( + full_steps if acceleration == ACCELERATION_NONE else accelerated_steps + ) + steps_high = _rounded_stage_share(strategy_budget, high_range_share) + steps_low = strategy_budget - steps_high + return steps_high, steps_low, strategy_budget, strategy_budget + + def _interpolate(values: Sequence[float], position: float) -> float: last = len(values) - 1 if position <= 0.0: @@ -394,6 +428,7 @@ def optimize_shift( def _warnings_for_plan( *, acceleration: str, + handoff_mode: str, steps: int, steps_high: int, steps_low: int, @@ -431,7 +466,22 @@ def _warnings_for_plan( elif not accelerate_low and steps_low < 3: warnings.append("The unaccelerated low-noise stage has very few steps.") + if handoff_mode == HANDOFF_PROGRESSIVE: + warnings.append( + "Progressive Transcode uses a clean high-stage output and re-noises " + "the re-encoded latent at the low-stage handoff sigma." + ) + if acceleration != ACCELERATION_NONE: + warnings.append( + "Progressive Transcode allocates the accelerated strategy budget " + "across both stages; the full-step budget is inactive for step " + "accounting in this mode." + ) + selected_budget = ( + steps + if handoff_mode == HANDOFF_PROGRESSIVE + else accelerated_steps if acceleration == ACCELERATION_BOTH else full_steps @@ -467,7 +517,8 @@ def _summary(plan: dict[str, Any]) -> str: if plan["warnings"]: warning_text = " | " + " ".join(plan["warnings"]) return ( - f"{plan['task']} · {plan['acceleration']} · {plan['priority']} | " + f"{plan['task']} · {plan['acceleration']} · {plan['priority']} · " + f"{plan['handoff_mode']} | " f"{plan['steps']} steps = {plan['steps_high']} high + " f"{plan['steps_low']} low | budgets A{plan['accelerated_steps']}/" f"F{plan['full_steps']} · range {plan['high_range_percent']:.1f}% high | " @@ -488,6 +539,7 @@ def build_plan( scheduler: str, priority: str, sigma_provider: SigmaProvider, + handoff_mode: str = HANDOFF_DIRECT, forced_steps_high: int | None = None, forced_shift: float | None = None, forced_high_range_percent: float | None = None, @@ -498,6 +550,8 @@ def build_plan( raise ValueError(f"Unknown Wan 2.2 acceleration mode: {acceleration}") if priority not in PRIORITY_OPTIONS: raise ValueError(f"Unknown Wan 2.2 priority: {priority}") + if handoff_mode not in HANDOFF_OPTIONS: + raise ValueError(f"Unknown Wan 2.2 handoff mode: {handoff_mode}") accelerated_steps = int(accelerated_steps) full_steps = int(full_steps) @@ -571,11 +625,12 @@ def build_plan( priority=priority, ) - steps_high, steps_low, high_budget, low_budget = project_stage_steps( + steps_high, steps_low, high_budget, low_budget = project_handoff_steps( accelerated_steps=accelerated_steps, full_steps=full_steps, acceleration=acceleration, high_range_share=high_range_share, + handoff_mode=handoff_mode, ) projected_steps_high = steps_high projected_steps_low = steps_low @@ -620,6 +675,18 @@ def build_plan( ) sigmas_high = sigmas[: steps_high + 1] sigmas_low = sigmas[steps_high:] + sigmas_high_terminal = _sequence_like( + sigmas_high, + [*sigma_values[:steps_high], 0.0], + ) + if handoff_mode == HANDOFF_PROGRESSIVE: + sigmas_high_sampling = sigmas_high_terminal + high_output_mode = "terminal" + low_noise_mode = "random" + else: + sigmas_high_sampling = sigmas_high + high_output_mode = "leftover" + low_noise_mode = "disable" boundary_sigma_before = sigma_values[steps_high - 1] boundary_sigma_after = sigma_values[steps_high] boundary_distance = min( @@ -637,6 +704,7 @@ def build_plan( warnings = _warnings_for_plan( acceleration=acceleration, + handoff_mode=handoff_mode, steps=steps, steps_high=steps_high, steps_low=steps_low, @@ -656,6 +724,7 @@ def build_plan( "full_steps": full_steps, "scheduler": scheduler, "priority": priority, + "handoff_mode": handoff_mode, "sigma_provider": sigma_provider, } plan: dict[str, Any] = { @@ -666,6 +735,7 @@ def build_plan( "overrides": overrides, "task": task, "acceleration": acceleration, + "handoff_mode": handoff_mode, "accelerate_high": acceleration in (ACCELERATION_HIGH, ACCELERATION_BOTH), "accelerate_low": acceleration @@ -703,6 +773,12 @@ def build_plan( "sigmas": sigmas, "sigmas_high": sigmas_high, "sigmas_low": sigmas_low, + "sigmas_high_terminal": sigmas_high_terminal, + "sigmas_high_sampling": sigmas_high_sampling, + "high_output_mode": high_output_mode, + "noise_high_mode": "random", + "low_noise_mode": low_noise_mode, + "noise_low_mode": low_noise_mode, "split_sigma": boundary_sigma_after, "boundary_sigma_before": boundary_sigma_before, "boundary_sigma_after": boundary_sigma_after, @@ -743,6 +819,7 @@ def validate_plan(plan: Any) -> dict[str, Any]: required = { "source_controls", "overrides", + "handoff_mode", "curve_profile", "curve_mode", "shift_anchor", @@ -754,6 +831,11 @@ def validate_plan(plan: Any) -> dict[str, Any]: "sigmas", "sigmas_high", "sigmas_low", + "sigmas_high_terminal", + "sigmas_high_sampling", + "high_output_mode", + "noise_high_mode", + "noise_low_mode", "boundary", } missing = sorted(required.difference(plan)) @@ -765,6 +847,8 @@ def validate_plan(plan: Any) -> dict[str, Any]: raise ValueError("Sampling Plan overrides must be a dictionary.") if plan["curve_mode"] not in {"exact", "piecewise"}: raise ValueError(f"Unknown Sampling Plan curve mode: {plan['curve_mode']}.") + if plan["handoff_mode"] not in HANDOFF_OPTIONS: + raise ValueError(f"Unknown Sampling Plan handoff mode: {plan['handoff_mode']}.") steps = int(plan["steps"]) steps_high = int(plan["steps_high"]) @@ -781,6 +865,15 @@ def validate_plan(plan: Any) -> dict[str, Any]: context="Sampling Plan high sigma curve", require_terminal_zero=False, ) + _, high_terminal_values = _validated_sigmas( + plan["sigmas_high_terminal"], + context="Sampling Plan terminal high sigma curve", + ) + _, high_sampling_values = _validated_sigmas( + plan["sigmas_high_sampling"], + context="Sampling Plan sampling high sigma curve", + require_terminal_zero=plan["handoff_mode"] == HANDOFF_PROGRESSIVE, + ) _, low_values = _validated_sigmas( plan["sigmas_low"], context="Sampling Plan low sigma curve", @@ -789,6 +882,14 @@ def validate_plan(plan: Any) -> dict[str, Any]: raise ValueError("Sampling Plan step count disagrees with its sigma curve.") if len(high_values) != steps_high + 1: raise ValueError("Sampling Plan high-step count disagrees with its sigma curve.") + if len(high_terminal_values) != steps_high + 1: + raise ValueError( + "Sampling Plan terminal high-step count disagrees with its sigma curve." + ) + if len(high_sampling_values) != steps_high + 1: + raise ValueError( + "Sampling Plan sampling high-step count disagrees with its sigma curve." + ) if len(low_values) != steps_low + 1: raise ValueError("Sampling Plan low-step count disagrees with its sigma curve.") if not _sequences_close(plan["sigmas_high"], plan["sigmas"][: steps_high + 1]): @@ -802,6 +903,27 @@ def validate_plan(plan: Any) -> dict[str, Any]: abs_tol=SIGMA_TOLERANCE, ): raise ValueError("Sampling Plan stages must share exactly one handoff sigma.") + if plan["handoff_mode"] == HANDOFF_PROGRESSIVE: + if not _sequences_close( + plan["sigmas_high_sampling"], + plan["sigmas_high_terminal"], + ): + raise ValueError( + "Progressive Transcode must use the terminal high sigma curve." + ) + if plan["noise_low_mode"] != "random" or plan["high_output_mode"] != "terminal": + raise ValueError( + "Progressive Transcode requires terminal high output and low re-noise." + ) + else: + if not _sequences_close(plan["sigmas_high_sampling"], plan["sigmas_high"]): + raise ValueError( + "Direct Latent must use the handoff high sigma curve." + ) + if plan["noise_low_mode"] != "disable" or plan["high_output_mode"] != "leftover": + raise ValueError( + "Direct Latent requires leftover high output and disabled low noise." + ) boundary = float(plan["boundary"]) before = values[steps_high - 1]