405 lines
16 KiB
Python
405 lines
16 KiB
Python
import math
|
|
import torch
|
|
import torch.nn.functional as F
|
|
import comfy.clip_vision
|
|
import comfy.model_management as mm
|
|
import comfy.utils
|
|
import node_helpers
|
|
|
|
|
|
def _normalize_image_output(decoded):
|
|
"""Convert decoded VAE output to ComfyUI's standard IMAGE shape."""
|
|
if decoded.ndim == 5 and decoded.shape[0] == 1:
|
|
return decoded[0]
|
|
return decoded
|
|
|
|
|
|
def _find_coprime_skip(n):
|
|
"""Find a shift_skip value > n//2 that is coprime with n, ensuring full position coverage."""
|
|
candidate = n // 2 + 1
|
|
while candidate < n:
|
|
if math.gcd(candidate, n) == 1:
|
|
return candidate
|
|
candidate += 1
|
|
return 1
|
|
|
|
|
|
class MobiusLatentShiftWrapper:
|
|
"""
|
|
Mobius latent shift adapted for WAN FLF conditioning.
|
|
|
|
Rolls the noisy latent each step so the model sees varying temporal RoPE positions.
|
|
Optional concat-conditioning rolling can be enabled experimentally when desired.
|
|
"""
|
|
|
|
def __init__(self, shift_skip, latent_frames, roll_concat_conditioning=False):
|
|
self.shift_skip = shift_skip
|
|
self.latent_frames = latent_frames
|
|
self.shift_idx = 0
|
|
self.active_shift = 0
|
|
self.last_sigma = None
|
|
self.roll_concat_conditioning = roll_concat_conditioning
|
|
|
|
def __call__(self, apply_model_func, args_dict):
|
|
sigma = args_dict["timestep"]
|
|
|
|
if self.last_sigma is None or not torch.equal(sigma, self.last_sigma):
|
|
self.last_sigma = sigma.clone()
|
|
self.active_shift = self.shift_idx
|
|
self.shift_idx = (self.shift_idx + self.shift_skip) % self.latent_frames
|
|
|
|
if self.active_shift == 0:
|
|
return apply_model_func(args_dict["input"], args_dict["timestep"], **args_dict["c"])
|
|
|
|
shift = -self.active_shift
|
|
x = torch.roll(args_dict["input"], shift, dims=2)
|
|
c = dict(args_dict["c"])
|
|
|
|
# Rolling FLF concat conditioning makes the fixed-frame constraints move
|
|
# through time, which can destabilize the tail of the clip.
|
|
if self.roll_concat_conditioning and "c_concat" in c and c["c_concat"] is not None:
|
|
c_concat = c["c_concat"]
|
|
if hasattr(c_concat, "cond"):
|
|
rolled = torch.roll(c_concat.cond, shift, dims=2)
|
|
c["c_concat"] = c_concat.__class__(rolled)
|
|
else:
|
|
c["c_concat"] = torch.roll(c_concat, shift, dims=2)
|
|
|
|
output = apply_model_func(x, args_dict["timestep"], **c)
|
|
return torch.roll(output, -shift, dims=2)
|
|
|
|
|
|
class LoopBoundaryBlend:
|
|
"""
|
|
Post-CFG callback that blends end-frame predictions toward start-frame predictions.
|
|
|
|
For a seamless loop (start == end image), the denoised predictions at the temporal
|
|
boundaries should converge. This corrects the known WAN end-frame noise artifact
|
|
by gradually replacing the last K frames' predictions with the first K frames'.
|
|
"""
|
|
|
|
def __init__(self, blend_frames, blend_strength):
|
|
self.blend_frames = blend_frames
|
|
self.blend_strength = blend_strength
|
|
|
|
def __call__(self, args):
|
|
denoised = args["denoised"]
|
|
T = denoised.shape[2]
|
|
K = min(self.blend_frames, T // 4)
|
|
|
|
if K <= 0 or self.blend_strength <= 0:
|
|
return denoised
|
|
|
|
ramp = torch.linspace(0.0, self.blend_strength, K, device=denoised.device, dtype=denoised.dtype)
|
|
ramp = ramp.view(1, 1, K, 1, 1)
|
|
|
|
result = denoised.clone()
|
|
result[:, :, -K:] = denoised[:, :, -K:] * (1.0 - ramp) + denoised[:, :, :K] * ramp
|
|
return result
|
|
|
|
|
|
class SmartMobiusWanLoop:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"model_high": ("MODEL",),
|
|
"model_low": ("MODEL",),
|
|
"positive": ("CONDITIONING",),
|
|
"negative": ("CONDITIONING",),
|
|
"vae": ("VAE",),
|
|
"width": ("INT", {"default": 832, "min": 16, "max": 8192, "step": 16}),
|
|
"height": ("INT", {"default": 480, "min": 16, "max": 8192, "step": 16}),
|
|
"length": ("INT", {"default": 81, "min": 1, "max": 8192, "step": 4}),
|
|
"batch_size": ("INT", {"default": 1, "min": 1, "max": 4096}),
|
|
},
|
|
"optional": {
|
|
"start_image": ("IMAGE",),
|
|
"end_image": ("IMAGE",),
|
|
"clip_vision_start_image": ("CLIP_VISION_OUTPUT",),
|
|
"clip_vision_end_image": ("CLIP_VISION_OUTPUT",),
|
|
"motion_amplitude": (
|
|
"FLOAT",
|
|
{"default": 1.0, "min": 1.0, "max": 2.0, "step": 0.05,
|
|
"tooltip": "1.0 = official FLF, 2.0 = max structural repulsion boost"},
|
|
),
|
|
"enable_latent_shift": (
|
|
"BOOLEAN",
|
|
{"default": True,
|
|
"tooltip": "Mobius-style circular latent rolling each denoising step"},
|
|
),
|
|
"shift_skip": (
|
|
"INT",
|
|
{"default": 0, "min": 0, "max": 100, "step": 1,
|
|
"tooltip": "Roll stride per step. 0 = auto (coprime with latent frame count)"},
|
|
),
|
|
"enable_loop_blend": (
|
|
"BOOLEAN",
|
|
{"default": True,
|
|
"tooltip": "Blend end-frame predictions toward start-frame predictions after CFG"},
|
|
),
|
|
"blend_frames": (
|
|
"INT",
|
|
{"default": 5, "min": 1, "max": 20, "step": 1,
|
|
"tooltip": "Number of end frames to correct via boundary blending"},
|
|
),
|
|
"blend_strength": (
|
|
"FLOAT",
|
|
{"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.05,
|
|
"tooltip": "0 = no correction, 1 = fully replace end with start predictions"},
|
|
),
|
|
"enable_circular_rope": (
|
|
"BOOLEAN",
|
|
{"default": True,
|
|
"tooltip": "Compress temporal RoPE so last frame is closer to first"},
|
|
),
|
|
"roll_concat_conditioning": (
|
|
"BOOLEAN",
|
|
{
|
|
"default": False,
|
|
"tooltip": "Experimental. Roll concat_mask/concat_latent_image together with the latent cycle."
|
|
},
|
|
),
|
|
"circular_rope_strength": (
|
|
"FLOAT",
|
|
{"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.05,
|
|
"tooltip": "0 = no change, 1 = maximum temporal compression"},
|
|
),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("MODEL", "MODEL", "CONDITIONING", "CONDITIONING", "LATENT")
|
|
RETURN_NAMES = ("model_high", "model_low", "positive", "negative", "latent")
|
|
FUNCTION = "execute"
|
|
CATEGORY = "SmartHelperNodes/Wan"
|
|
DESCRIPTION = (
|
|
"Mobius seamless loop for WAN 2.2 dual-pass (high/low noise). "
|
|
"Three independent mechanisms: (1) Latent shift — cycles frame positions across "
|
|
"denoising steps to distribute RoPE bias. (2) Loop blend — corrects end-frame noise "
|
|
"by blending toward start-frame predictions. (3) Circular RoPE — compresses temporal "
|
|
"positions so the last frame is closer to the first in attention space."
|
|
)
|
|
|
|
def execute(
|
|
self,
|
|
model_high,
|
|
model_low,
|
|
positive,
|
|
negative,
|
|
vae,
|
|
width,
|
|
height,
|
|
length,
|
|
batch_size,
|
|
start_image=None,
|
|
end_image=None,
|
|
clip_vision_start_image=None,
|
|
clip_vision_end_image=None,
|
|
motion_amplitude=1.0,
|
|
enable_latent_shift=True,
|
|
shift_skip=0,
|
|
enable_loop_blend=True,
|
|
blend_frames=5,
|
|
blend_strength=0.5,
|
|
enable_circular_rope=True,
|
|
roll_concat_conditioning=False,
|
|
circular_rope_strength=0.5,
|
|
):
|
|
spacial_scale = vae.spacial_compression_encode()
|
|
latent_frames = ((length - 1) // 4) + 1
|
|
|
|
latent = torch.zeros(
|
|
[batch_size, vae.latent_channels, latent_frames, height // spacial_scale, width // spacial_scale],
|
|
device=mm.intermediate_device(),
|
|
)
|
|
|
|
# --- Image preprocessing ---
|
|
if start_image is not None:
|
|
start_image = comfy.utils.common_upscale(
|
|
start_image[:length].movedim(-1, 1), width, height, "bilinear", "center"
|
|
).movedim(1, -1)
|
|
if end_image is not None:
|
|
end_image = comfy.utils.common_upscale(
|
|
end_image[-length:].movedim(-1, 1), width, height, "bilinear", "center"
|
|
).movedim(1, -1)
|
|
|
|
# --- Build reference image sequence (gray fill + start/end frames) ---
|
|
official_image = torch.ones((length, height, width, 3), device=mm.intermediate_device()) * 0.5
|
|
mask = torch.ones(
|
|
(1, 1, latent_frames * 4, height // spacial_scale, width // spacial_scale),
|
|
device=mm.intermediate_device(),
|
|
)
|
|
|
|
if start_image is not None:
|
|
official_image[:start_image.shape[0]] = start_image
|
|
mask[:, :, :start_image.shape[0] + 3] = 0.0
|
|
if end_image is not None:
|
|
official_image[-end_image.shape[0]:] = end_image
|
|
mask[:, :, -end_image.shape[0]:] = 0.0
|
|
|
|
official_latent = vae.encode(official_image[:, :, :, :3])
|
|
|
|
# --- Inverse structural repulsion (PainterFLF2V algorithm) ---
|
|
if length > 2 and motion_amplitude > 1.001 and start_image is not None and end_image is not None:
|
|
start_l = official_latent[:, :, 0:1]
|
|
end_l = official_latent[:, :, -1:]
|
|
t = torch.linspace(0.0, 1.0, official_latent.shape[2], device=official_latent.device).view(1, 1, -1, 1, 1)
|
|
linear_latent = start_l * (1 - t) + end_l * t
|
|
|
|
diff = official_latent - linear_latent
|
|
h, w = diff.shape[-2], diff.shape[-1]
|
|
low_freq_diff = F.interpolate(
|
|
diff.view(-1, vae.latent_channels, h, w), size=(max(1, h // 8), max(1, w // 8)), mode="area"
|
|
)
|
|
low_freq_diff = F.interpolate(low_freq_diff, size=(h, w), mode="bilinear")
|
|
low_freq_diff = low_freq_diff.view_as(diff)
|
|
high_freq_diff = diff - low_freq_diff
|
|
|
|
boost_scale = (motion_amplitude - 1.0) * 4.0
|
|
concat_latent_image = official_latent + (high_freq_diff * boost_scale)
|
|
else:
|
|
concat_latent_image = official_latent
|
|
|
|
# --- Mask reshape: [1,1,T*4,H,W] -> [1,4,T,H,W] ---
|
|
mask = mask.view(1, mask.shape[2] // 4, 4, mask.shape[3], mask.shape[4]).transpose(1, 2)
|
|
|
|
# --- Inject conditioning ---
|
|
positive = node_helpers.conditioning_set_values(
|
|
positive, {"concat_latent_image": concat_latent_image, "concat_mask": mask}
|
|
)
|
|
negative = node_helpers.conditioning_set_values(
|
|
negative, {"concat_latent_image": concat_latent_image, "concat_mask": mask}
|
|
)
|
|
|
|
# --- CLIP vision merge ---
|
|
clip_vision_output = None
|
|
if clip_vision_start_image is not None:
|
|
clip_vision_output = clip_vision_start_image
|
|
if clip_vision_end_image is not None:
|
|
if clip_vision_output is not None:
|
|
states = torch.cat(
|
|
[clip_vision_output.penultimate_hidden_states,
|
|
clip_vision_end_image.penultimate_hidden_states],
|
|
dim=-2,
|
|
)
|
|
clip_vision_output = comfy.clip_vision.Output()
|
|
clip_vision_output.penultimate_hidden_states = states
|
|
else:
|
|
clip_vision_output = clip_vision_end_image
|
|
if clip_vision_output is not None:
|
|
positive = node_helpers.conditioning_set_values(positive, {"clip_vision_output": clip_vision_output})
|
|
negative = node_helpers.conditioning_set_values(negative, {"clip_vision_output": clip_vision_output})
|
|
|
|
# --- Model patching (shared state across both models) ---
|
|
patched_high = model_high.clone()
|
|
patched_low = model_low.clone()
|
|
|
|
if enable_latent_shift:
|
|
effective_skip = shift_skip if shift_skip > 0 else _find_coprime_skip(latent_frames)
|
|
if math.gcd(effective_skip, latent_frames) != 1:
|
|
effective_skip = _find_coprime_skip(latent_frames)
|
|
wrapper = MobiusLatentShiftWrapper(
|
|
effective_skip,
|
|
latent_frames,
|
|
roll_concat_conditioning=roll_concat_conditioning,
|
|
)
|
|
patched_high.set_model_unet_function_wrapper(wrapper)
|
|
patched_low.set_model_unet_function_wrapper(wrapper)
|
|
|
|
if enable_loop_blend and blend_strength > 0 and blend_frames > 0:
|
|
blend_fn = LoopBoundaryBlend(blend_frames, blend_strength)
|
|
patched_high.set_model_sampler_post_cfg_function(blend_fn)
|
|
patched_low.set_model_sampler_post_cfg_function(blend_fn)
|
|
|
|
if enable_circular_rope and circular_rope_strength > 0.0:
|
|
scale_t = 1.0 - (circular_rope_strength / latent_frames)
|
|
for m in (patched_high, patched_low):
|
|
to = m.model_options.get("transformer_options", {})
|
|
existing_rope = to.get("rope_options", {})
|
|
existing_rope["scale_t"] = scale_t
|
|
to["rope_options"] = existing_rope
|
|
m.model_options["transformer_options"] = to
|
|
|
|
out_latent = {"samples": latent}
|
|
return (patched_high, patched_low, positive, negative, out_latent)
|
|
|
|
|
|
class SmartMobiusVAEDecode:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"samples": ("LATENT",),
|
|
"vae": ("VAE",),
|
|
},
|
|
"optional": {
|
|
"enable_frame_invariance": (
|
|
"BOOLEAN",
|
|
{"default": True, "tooltip": "Decode boundary region with wrapped context to fix 3D VAE seam artifacts"},
|
|
),
|
|
"boundary_frames": (
|
|
"INT",
|
|
{"default": 3, "min": 2, "max": 5, "step": 1,
|
|
"tooltip": "Latent frames to wrap at each boundary side"},
|
|
),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("image",)
|
|
FUNCTION = "execute"
|
|
CATEGORY = "SmartHelperNodes/Wan"
|
|
DESCRIPTION = (
|
|
"VAE decode with optional frame-invariance for seamless loops. "
|
|
"Wraps boundary latent frames and decodes them with full temporal context "
|
|
"to eliminate 3D VAE seam artifacts at the loop point."
|
|
)
|
|
|
|
def execute(self, samples, vae, enable_frame_invariance=True, boundary_frames=3):
|
|
latent = samples["samples"]
|
|
|
|
if not enable_frame_invariance or latent.shape[2] < boundary_frames * 2 + 1:
|
|
decoded = _normalize_image_output(vae.decode(latent))
|
|
return (decoded,)
|
|
|
|
k = boundary_frames
|
|
|
|
main_pixels = vae.decode(latent)
|
|
|
|
boundary_latent = torch.cat([latent[:, :, -k:], latent[:, :, :k]], dim=2)
|
|
boundary_pixels = vae.decode(boundary_latent)
|
|
frame_dim = 1 if boundary_pixels.ndim == 5 else 0
|
|
|
|
left_pixel_count = 1 + (k - 1) * 4
|
|
tail_slice = [slice(None)] * boundary_pixels.ndim
|
|
tail_slice[frame_dim] = slice(0, left_pixel_count)
|
|
tail_good = boundary_pixels[tuple(tail_slice)]
|
|
|
|
head_slice = [slice(None)] * boundary_pixels.ndim
|
|
head_slice[frame_dim] = slice(left_pixel_count, None)
|
|
head_good = boundary_pixels[tuple(head_slice)]
|
|
|
|
result = main_pixels.clone()
|
|
frame_dim = 1 if result.ndim == 5 else 0
|
|
|
|
tail_slice = [slice(None)] * result.ndim
|
|
tail_slice[frame_dim] = slice(-tail_good.shape[frame_dim], None)
|
|
result[tuple(tail_slice)] = tail_good
|
|
|
|
head_slice = [slice(None)] * result.ndim
|
|
head_slice[frame_dim] = slice(0, head_good.shape[frame_dim])
|
|
result[tuple(head_slice)] = head_good
|
|
|
|
return (_normalize_image_output(result),)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"SmartMobiusWanLoop": SmartMobiusWanLoop,
|
|
"SmartMobiusVAEDecode": SmartMobiusVAEDecode,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"SmartMobiusWanLoop": "Smart Mobius WAN Loop",
|
|
"SmartMobiusVAEDecode": "Smart Mobius VAE Decode",
|
|
}
|