From 288a7e96a97a8a233cd1a8b993adc28c640bca70 Mon Sep 17 00:00:00 2001
From: kijai <40791699+kijai@users.noreply.github.com>
Date: Mon, 25 Aug 2025 14:53:36 +0300
Subject: [PATCH] Add experimental scheduler node with sigma plot
---
.../wanvideo_FLF2V_720P_example_02.json | 421 ++++++++++--------
nodes.py | 88 +++-
wanvideo/schedulers/__init__.py | 12 +-
3 files changed, 313 insertions(+), 208 deletions(-)
diff --git a/example_workflows/wanvideo_FLF2V_720P_example_02.json b/example_workflows/wanvideo_FLF2V_720P_example_02.json
index 8d37929..a2fd652 100644
--- a/example_workflows/wanvideo_FLF2V_720P_example_02.json
+++ b/example_workflows/wanvideo_FLF2V_720P_example_02.json
@@ -611,6 +611,7 @@
},
{
"name": "image_2",
+ "shape": 7,
"type": "IMAGE",
"link": 152
},
@@ -662,6 +663,7 @@
},
{
"name": "image_2",
+ "shape": 7,
"type": "IMAGE",
"link": 150
}
@@ -806,7 +808,14 @@
"flags": {},
"order": 9,
"mode": 0,
- "inputs": [],
+ "inputs": [
+ {
+ "name": "compile_args",
+ "shape": 7,
+ "type": "WANCOMPILEARGS",
+ "link": null
+ }
+ ],
"outputs": [
{
"name": "vae",
@@ -903,7 +912,7 @@
],
"size": [
887.1368408203125,
- 934.646484375
+ 334
],
"flags": {},
"order": 38,
@@ -988,6 +997,7 @@
"inputs": [
{
"name": "vae",
+ "shape": 7,
"type": "WANVAE",
"link": 170
},
@@ -1095,6 +1105,7 @@
"inputs": [
{
"name": "t5",
+ "shape": 7,
"type": "WANTEXTENCODER",
"link": 15
},
@@ -1123,7 +1134,9 @@
"widgets_values": [
"CG动画风格，一只蓝色的小鸟从地面起飞，煽动翅膀。小鸟羽毛细腻，胸前有独特的花纹，背景是蓝天白云，阳光明媚。镜跟随小鸟向上移动，展现出小鸟飞翔的姿态和天空的广阔。近景，仰视视角",
"色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走",
- true
+ true,
+ false,
+ "gpu"
],
"color": "#332922",
"bgcolor": "#593930"
@@ -1219,7 +1232,7 @@
],
"size": [
315,
- 154
+ 202
],
"flags": {},
"order": 11,
@@ -1245,7 +1258,9 @@
false,
false,
true,
- 0
+ 0,
+ 0,
+ false
],
"color": "#223",
"bgcolor": "#335"
@@ -1477,7 +1492,8 @@
"0, 0, 0",
"center",
16,
- "cpu"
+ "cpu",
+ "| Output: | 1 x 640 x 640 | 4.69MB |
"
]
},
{
@@ -1489,7 +1505,7 @@
],
"size": [
270,
- 286
+ 336
],
"flags": {},
"order": 28,
@@ -1564,9 +1580,199 @@
"0, 0, 0",
"center",
16,
- "cpu"
+ "cpu",
+ "| Output: | 1 x 640 x 640 | 4.69MB |
"
]
},
+ {
+ "id": 106,
+ "type": "WanVideoLoraSelect",
+ "pos": [
+ -336.7720642089844,
+ -698.3348999023438
+ ],
+ "size": [
+ 424.9496765136719,
+ 150
+ ],
+ "flags": {},
+ "order": 17,
+ "mode": 0,
+ "inputs": [
+ {
+ "name": "prev_lora",
+ "shape": 7,
+ "type": "WANVIDLORA",
+ "link": null
+ },
+ {
+ "name": "blocks",
+ "shape": 7,
+ "type": "SELECTEDBLOCKS",
+ "link": null
+ }
+ ],
+ "outputs": [
+ {
+ "name": "lora",
+ "type": "WANVIDLORA",
+ "links": [
+ 179
+ ]
+ }
+ ],
+ "properties": {
+ "cnr_id": "ComfyUI-WanVideoWrapper",
+ "ver": "974dd656dab305f7fa122cca435759105ea44488",
+ "Node name for S&R": "WanVideoLoraSelect"
+ },
+ "widgets_values": [
+ "Wan21_T2V_14B_lightx2v_cfg_step_distill_lora_rank32.safetensors",
+ 1.2000000000000002,
+ false,
+ true
+ ],
+ "color": "#223",
+ "bgcolor": "#335"
+ },
+ {
+ "id": 35,
+ "type": "WanVideoTorchCompileSettings",
+ "pos": [
+ -307.4797058105469,
+ -1197.4749755859375
+ ],
+ "size": [
+ 421.6000061035156,
+ 202
+ ],
+ "flags": {},
+ "order": 18,
+ "mode": 0,
+ "inputs": [],
+ "outputs": [
+ {
+ "name": "torch_compile_args",
+ "type": "WANCOMPILEARGS",
+ "slot_index": 0,
+ "links": [
+ 190
+ ]
+ }
+ ],
+ "properties": {
+ "cnr_id": "ComfyUI-WanVideoWrapper",
+ "ver": "d9b1f4d1a5aea91d101ae97a54714a5861af3f50",
+ "Node name for S&R": "WanVideoTorchCompileSettings"
+ },
+ "widgets_values": [
+ "inductor",
+ false,
+ "default",
+ false,
+ 64,
+ true,
+ 128
+ ],
+ "color": "#223",
+ "bgcolor": "#335"
+ },
+ {
+ "id": 22,
+ "type": "WanVideoModelLoader",
+ "pos": [
+ 119.37029266357422,
+ -926.8419799804688
+ ],
+ "size": [
+ 477.4410095214844,
+ 314
+ ],
+ "flags": {},
+ "order": 24,
+ "mode": 0,
+ "inputs": [
+ {
+ "name": "compile_args",
+ "shape": 7,
+ "type": "WANCOMPILEARGS",
+ "link": 190
+ },
+ {
+ "name": "block_swap_args",
+ "shape": 7,
+ "type": "BLOCKSWAPARGS",
+ "link": 174
+ },
+ {
+ "name": "lora",
+ "shape": 7,
+ "type": "WANVIDLORA",
+ "link": 179
+ },
+ {
+ "name": "vram_management_args",
+ "shape": 7,
+ "type": "VRAM_MANAGEMENTARGS",
+ "link": null
+ },
+ {
+ "name": "extra_model",
+ "shape": 7,
+ "type": "VACEPATH",
+ "link": null
+ },
+ {
+ "name": "fantasytalking_model",
+ "shape": 7,
+ "type": "FANTASYTALKINGMODEL",
+ "link": null
+ },
+ {
+ "name": "multitalk_model",
+ "shape": 7,
+ "type": "MULTITALKMODEL",
+ "link": null
+ },
+ {
+ "name": "fantasyportrait_model",
+ "shape": 7,
+ "type": "FANTASYPORTRAITMODEL",
+ "link": null
+ },
+ {
+ "name": "vace_model",
+ "shape": 7,
+ "type": "VACEPATH",
+ "link": null
+ }
+ ],
+ "outputs": [
+ {
+ "name": "model",
+ "type": "WANVIDEOMODEL",
+ "slot_index": 0,
+ "links": [
+ 29,
+ 103
+ ]
+ }
+ ],
+ "properties": {
+ "cnr_id": "ComfyUI-WanVideoWrapper",
+ "ver": "d9b1f4d1a5aea91d101ae97a54714a5861af3f50",
+ "Node name for S&R": "WanVideoModelLoader"
+ },
+ "widgets_values": [
+ "WanVideo\\Wan2_1-FLF2V-14B-720P_fp8_e4m3fn.safetensors",
+ "fp16_fast",
+ "fp8_e4m3fn",
+ "offload_device",
+ "sageattn"
+ ],
+ "color": "#223",
+ "bgcolor": "#335"
+ },
{
"id": 27,
"type": "WanVideoSampler",
@@ -1675,6 +1881,12 @@
"shape": 7,
"type": "MULTITALK_EMBEDS",
"link": null
+ },
+ {
+ "name": "freeinit_args",
+ "shape": 7,
+ "type": "FREEINITARGS",
+ "link": null
}
],
"outputs": [
@@ -1685,6 +1897,11 @@
"links": [
166
]
+ },
+ {
+ "name": "denoised_samples",
+ "type": "LATENT",
+ "links": null
}
],
"properties": {
@@ -1704,184 +1921,10 @@
1,
"",
"comfy",
- ""
- ]
- },
- {
- "id": 106,
- "type": "WanVideoLoraSelect",
- "pos": [
- -336.7720642089844,
- -698.3348999023438
- ],
- "size": [
- 424.9496765136719,
- 126
- ],
- "flags": {},
- "order": 17,
- "mode": 0,
- "inputs": [
- {
- "name": "prev_lora",
- "shape": 7,
- "type": "WANVIDLORA",
- "link": null
- },
- {
- "name": "blocks",
- "shape": 7,
- "type": "SELECTEDBLOCKS",
- "link": null
- }
- ],
- "outputs": [
- {
- "name": "lora",
- "type": "WANVIDLORA",
- "links": [
- 179
- ]
- }
- ],
- "properties": {
- "cnr_id": "ComfyUI-WanVideoWrapper",
- "ver": "974dd656dab305f7fa122cca435759105ea44488",
- "Node name for S&R": "WanVideoLoraSelect"
- },
- "widgets_values": [
- "Wan21_T2V_14B_lightx2v_cfg_step_distill_lora_rank32.safetensors",
- 1.2000000000000002,
+ 0,
+ -1,
false
- ],
- "color": "#223",
- "bgcolor": "#335"
- },
- {
- "id": 35,
- "type": "WanVideoTorchCompileSettings",
- "pos": [
- -307.4797058105469,
- -1197.4749755859375
- ],
- "size": [
- 421.6000061035156,
- 202
- ],
- "flags": {},
- "order": 18,
- "mode": 0,
- "inputs": [],
- "outputs": [
- {
- "name": "torch_compile_args",
- "type": "WANCOMPILEARGS",
- "slot_index": 0,
- "links": [
- 190
- ]
- }
- ],
- "properties": {
- "cnr_id": "ComfyUI-WanVideoWrapper",
- "ver": "d9b1f4d1a5aea91d101ae97a54714a5861af3f50",
- "Node name for S&R": "WanVideoTorchCompileSettings"
- },
- "widgets_values": [
- "inductor",
- false,
- "default",
- false,
- 64,
- true,
- 128
- ],
- "color": "#223",
- "bgcolor": "#335"
- },
- {
- "id": 22,
- "type": "WanVideoModelLoader",
- "pos": [
- 119.37029266357422,
- -926.8419799804688
- ],
- "size": [
- 477.4410095214844,
- 274
- ],
- "flags": {},
- "order": 24,
- "mode": 0,
- "inputs": [
- {
- "name": "compile_args",
- "shape": 7,
- "type": "WANCOMPILEARGS",
- "link": 190
- },
- {
- "name": "block_swap_args",
- "shape": 7,
- "type": "BLOCKSWAPARGS",
- "link": 174
- },
- {
- "name": "lora",
- "shape": 7,
- "type": "WANVIDLORA",
- "link": 179
- },
- {
- "name": "vram_management_args",
- "shape": 7,
- "type": "VRAM_MANAGEMENTARGS",
- "link": null
- },
- {
- "name": "vace_model",
- "shape": 7,
- "type": "VACEPATH",
- "link": null
- },
- {
- "name": "fantasytalking_model",
- "shape": 7,
- "type": "FANTASYTALKINGMODEL",
- "link": null
- },
- {
- "name": "multitalk_model",
- "shape": 7,
- "type": "MULTITALKMODEL",
- "link": null
- }
- ],
- "outputs": [
- {
- "name": "model",
- "type": "WANVIDEOMODEL",
- "slot_index": 0,
- "links": [
- 29,
- 103
- ]
- }
- ],
- "properties": {
- "cnr_id": "ComfyUI-WanVideoWrapper",
- "ver": "d9b1f4d1a5aea91d101ae97a54714a5861af3f50",
- "Node name for S&R": "WanVideoModelLoader"
- },
- "widgets_values": [
- "WanVideo\\Wan2_1-FLF2V-14B-720P_fp8_e4m3fn.safetensors",
- "fp16_fast",
- "fp8_e4m3fn",
- "offload_device",
- "sageattn"
- ],
- "color": "#223",
- "bgcolor": "#335"
+ ]
}
],
"links": [
@@ -2216,13 +2259,13 @@
"config": {},
"extra": {
"ds": {
- "scale": 0.6727499949326076,
+ "scale": 0.6115909044841886,
"offset": [
- 359.0502881120043,
- 1009.7385911003805
+ 710.2610924809787,
+ 967.8431584929548
]
},
- "frontendVersion": "1.23.4",
+ "frontendVersion": "1.26.3",
"node_versions": {
"ComfyUI-WanVideoWrapper": "f8f423eceeadf2edcb58fab73701333e83ca733e",
"comfy-core": "0.3.26",
diff --git a/nodes.py b/nodes.py
index 1c44246..422201b 100644
--- a/nodes.py
+++ b/nodes.py
@@ -1582,17 +1582,78 @@ class WanVideoScheduler: #WIP
def INPUT_TYPES(s):
return {"required": {
"scheduler": (scheduler_list, {"default": "unipc"}),
+ "steps": ("INT", {"default": 30, "min": 1, "tooltip": "Number of steps for the scheduler"}),
+ "shift": ("FLOAT", {"default": 5.0, "min": 0.0, "max": 1000.0, "step": 0.01}),
+ "start_step": ("INT", {"default": 0, "min": 0, "tooltip": "Starting step for the scheduler"}),
+ "end_step": ("INT", {"default": -1, "min": -1, "tooltip": "Ending step for the scheduler"})
+ },
+ "optional": {
+ "sigmas": ("SIGMAS", ),
+ },
+ "hidden": {
+ "unique_id": "UNIQUE_ID",
},
}
- RETURN_TYPES = (scheduler_list, )
- RETURN_NAMES = ("scheduler",)
+ RETURN_TYPES = ("SIGMAS", "INT", "FLOAT", scheduler_list, "INT", "INT",)
+ RETURN_NAMES = ("sigmas", "steps", "shift", "scheduler", "start_step", "end_step")
FUNCTION = "process"
CATEGORY = "WanVideoWrapper"
EXPERIMENTAL = True
- def process(self, scheduler):
- return (scheduler,)
+ def process(self, scheduler, steps, start_step, end_step, shift, unique_id, sigmas=None):
+ sample_scheduler, timesteps = get_scheduler(
+ scheduler,
+ steps,
+ start_step, end_step, shift,
+ device,
+ sigmas=sigmas)
+
+ scheduler_dict = {
+ "sample_scheduler": sample_scheduler,
+ "timesteps": timesteps,
+ }
+
+ try:
+ from server import PromptServer
+ import io
+ import base64
+ import matplotlib.pyplot as plt
+ except:
+ PromptServer = None
+ if unique_id and PromptServer is not None:
+ try:
+ # Plot sigmas and save to a buffer
+ sigmas_np = sample_scheduler.full_sigmas[:-1].cpu().numpy()
+ buf = io.BytesIO()
+ fig = plt.figure(facecolor='#353535')
+ ax = fig.add_subplot(111)
+ ax.set_facecolor('#353535') # Set axes background color
+ ax.plot(sigmas_np)
+ ax.set_title("Sigmas", color='white') # Title font color
+ ax.set_xlabel("Step", color='white') # X label font color
+ ax.set_ylabel("Sigma Value", color='white') # Y label font color
+ ax.tick_params(axis='x', colors='white') # X tick color
+ ax.tick_params(axis='y', colors='white') # Y tick color
+ # Add split point if end_step is defined
+ if end_step != -1 and 0 <= end_step < len(sigmas_np):
+ ax.axvline(end_step, color='red', linestyle='--', linewidth=2, label='end_step split')
+ ax.legend()
+ plt.tight_layout()
+ plt.savefig(buf, format='png')
+ plt.close(fig)
+ buf.seek(0)
+ img_base64 = base64.b64encode(buf.read()).decode('utf-8')
+ buf.close()
+
+ # Send as HTML img tag with base64 data
+ html_img = f"
"
+ PromptServer.instance.send_progress_text(html_img, unique_id)
+ except Exception as e:
+ print("Failed to send sigmas plot:", e)
+ pass
+
+ return (sigmas, steps, shift, scheduler_dict, start_step, end_step)
rope_functions = ["default", "comfy", "comfy_chunked"]
class WanVideoRoPEFunction:
@@ -1744,8 +1805,11 @@ class WanVideoSampler:
#region Scheduler
sample_scheduler = None
- if scheduler != "multitalk":
- sample_scheduler, timesteps, scheduler_step_args = get_scheduler(scheduler, steps, start_step, end_step, shift, device, transformer.dim, flowedit_args, denoise_strength, sigmas=sigmas, seed_g=seed_g)
+ if isinstance(scheduler, dict):
+ sample_scheduler = scheduler["sample_scheduler"]
+ timesteps = scheduler["timesteps"]
+ elif scheduler != "multitalk":
+ sample_scheduler, timesteps = get_scheduler(scheduler, steps, start_step, end_step, shift, device, transformer.dim, flowedit_args, denoise_strength, sigmas=sigmas)
log.info(f"sigmas: {sample_scheduler.sigmas}")
else:
timesteps = torch.tensor([1000, 750, 500, 250], device=device)
@@ -1761,7 +1825,11 @@ class WanVideoSampler:
start_step = steps - int(steps * denoise_strength) - 1
add_noise_to_samples = True #for now to not break old workflows
- noise_pred_flipped = None
+ scheduler_step_args = {"generator": seed_g}
+ step_sig = inspect.signature(sample_scheduler.step)
+ for arg in list(scheduler_step_args.keys()):
+ if arg not in step_sig.parameters:
+ scheduler_step_args.pop(arg)
if isinstance(cfg, list):
if steps < len(cfg):
@@ -1778,7 +1846,7 @@ class WanVideoSampler:
vace_data = vace_context = vace_scale = None
fun_or_fl2v_model = has_ref = drop_last = False
phantom_latents = fun_ref_image = ATI_tracks = None
- add_cond = attn_cond = attn_cond_neg = None
+ add_cond = attn_cond = attn_cond_neg = noise_pred_flipped = None
#I2V
image_cond = image_embeds.get("image_embeds", None)
@@ -2792,7 +2860,7 @@ class WanVideoSampler:
# FreeInit noise reinitialization (after first iteration)
if freeinit_args is not None and iter_idx > 0:
# restart scheduler for each iteration
- sample_scheduler, timesteps, scheduler_step_args = get_scheduler(scheduler, steps, start_step, end_step, shift, device, transformer.dim, flowedit_args, denoise_strength, sigmas=sigmas, seed_g=seed_g)
+ sample_scheduler, timesteps = get_scheduler(scheduler, steps, start_step, end_step, shift, device, transformer.dim, flowedit_args, denoise_strength, sigmas=sigmas)
# Re-apply start_step and end_step logic to timesteps and sigmas
if end_step != -1:
@@ -3310,7 +3378,7 @@ class WanVideoSampler:
timesteps = [torch.tensor([t], device=device) for t in timesteps]
timesteps = [timestep_transform(t, shift=shift, num_timesteps=1000) for t in timesteps]
else:
- sample_scheduler, timesteps, scheduler_step_args = get_scheduler(scheduler, total_steps, start_step, end_step, shift, device, transformer.dim, flowedit_args, denoise_strength, sigmas=sigmas, seed_g=seed_g)
+ sample_scheduler, timesteps = get_scheduler(scheduler, total_steps, start_step, end_step, shift, device, transformer.dim, flowedit_args, denoise_strength, sigmas=sigmas)
timesteps = [torch.tensor([float(t)], device=device) for t in timesteps] + [torch.tensor([0.], device=device)]
# sample videos
diff --git a/wanvideo/schedulers/__init__.py b/wanvideo/schedulers/__init__.py
index dc63457..407e107 100644
--- a/wanvideo/schedulers/__init__.py
+++ b/wanvideo/schedulers/__init__.py
@@ -23,7 +23,7 @@ scheduler_list = [
"multitalk"
]
-def get_scheduler(scheduler, steps, start_step, end_step, shift, device, transformer_dim, flowedit_args, denoise_strength, sigmas=None, seed_g=None):
+def get_scheduler(scheduler, steps, start_step, end_step, shift, device, transformer_dim=5120, flowedit_args=None, denoise_strength=1.0, sigmas=None):
timesteps = None
if 'unipc' in scheduler:
sample_scheduler = FlowUniPCMultistepScheduler(shift=shift)
@@ -130,6 +130,7 @@ def get_scheduler(scheduler, steps, start_step, end_step, shift, device, transfo
# Slice timesteps and sigmas once, based on indices
timesteps = timesteps[start_idx:end_idx+1]
+ sample_scheduler.full_sigmas = sample_scheduler.sigmas.clone()
sample_scheduler.sigmas = sample_scheduler.sigmas[start_idx:start_idx+len(timesteps)+1] # always one longer
@@ -138,11 +139,4 @@ def get_scheduler(scheduler, steps, start_step, end_step, shift, device, transfo
if hasattr(sample_scheduler, 'timesteps'):
sample_scheduler.timesteps = timesteps
- if seed_g is not None:
- scheduler_step_args = {"generator": seed_g}
- step_sig = inspect.signature(sample_scheduler.step)
- for arg in list(scheduler_step_args.keys()):
- if arg not in step_sig.parameters:
- scheduler_step_args.pop(arg)
-
- return sample_scheduler, timesteps, scheduler_step_args
\ No newline at end of file
+ return sample_scheduler, timesteps
\ No newline at end of file