nodes/boobkake22-ComfyUI-SamplingPlanner
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Add progressive handoff sampling mode

Browse Source
This commit is contained in:
Boobkake22
2026-06-27 15:09:06 -07:00
parent baa4a3776f
commit e8d8fcc81d
5 changed files with 436 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+68 -6
View File
@@ -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.
nodes.py
+145 -1
View File
@@ -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",
}
tests/test_nodes_smoke.py
+69
View File
@@ -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"]
tests/test_wan22_planner.py
+29
View File
@@ -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(
wan22_planner.py
+125 -3
View File
@@ -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]