310 lines
15 KiB
Python
310 lines
15 KiB
Python
from typing_extensions import override
|
|
from comfy_api.latest import io
|
|
import torch
|
|
import torch.nn.functional as F
|
|
import json
|
|
from typing import List, Tuple, Optional, Any
|
|
import node_helpers
|
|
import comfy
|
|
import comfy.utils
|
|
|
|
|
|
class WanMultiFrameRefToVideo(io.ComfyNode):
|
|
|
|
RETURN_TYPES = ("CONDITIONING", "CONDITIONING", "CONDITIONING", "LATENT")
|
|
RETURN_NAMES = ("positive_high", "positive_low", "negative", "latent")
|
|
CATEGORY = "ComfyUI-Wan22FMLF"
|
|
FUNCTION = "execute"
|
|
OUTPUT_NODE = False
|
|
|
|
@classmethod
|
|
def define_schema(cls) -> io.Schema:
|
|
return io.Schema(
|
|
node_id="WanMultiFrameRefToVideo",
|
|
display_name="Wan Multi-Frame Reference",
|
|
category="ComfyUI-Wan22FMLF",
|
|
inputs=[
|
|
io.Conditioning.Input("positive"),
|
|
io.Conditioning.Input("negative"),
|
|
io.Vae.Input("vae"),
|
|
io.Int.Input("width", default=832, min=16, max=8192, step=16, display_mode=io.NumberDisplay.number),
|
|
io.Int.Input("height", default=480, min=16, max=8192, step=16, display_mode=io.NumberDisplay.number),
|
|
io.Int.Input("length", default=81, min=1, max=8192, step=4, display_mode=io.NumberDisplay.number),
|
|
io.Int.Input("batch_size", default=1, min=1, max=4096, display_mode=io.NumberDisplay.number),
|
|
io.Image.Input("ref_images"),
|
|
io.Combo.Input("mode", ["NORMAL", "SINGLE_PERSON"], default="NORMAL", optional=True),
|
|
io.String.Input("ref_positions", default="", optional=True),
|
|
io.Float.Input("ref_strength_high", default=0.8, min=0.0, max=1.0, step=0.05, round=0.01, display_mode=io.NumberDisplay.slider, optional=True, display_name="ref_strength_high"),
|
|
io.Float.Input("end_frame_strength_high", default=1.0, min=0.0, max=1.0, step=0.05, round=0.01, display_mode=io.NumberDisplay.slider, optional=True, display_name="end_frame_strength_high"),
|
|
io.Float.Input("ref_strength_low", default=0.2, min=0.0, max=1.0, step=0.05, round=0.01, display_mode=io.NumberDisplay.slider, optional=True, display_name="ref_strength_low"),
|
|
io.Float.Input("start_frame_strength_low", default=1.0, min=0.0, max=1.0, step=0.05, round=0.01, display_mode=io.NumberDisplay.slider, optional=True, display_name="start_frame_strength_low"),
|
|
io.Float.Input("end_frame_strength_low", default=1.0, min=0.0, max=1.0, step=0.05, round=0.01, display_mode=io.NumberDisplay.slider, optional=True, display_name="end_frame_strength_low"),
|
|
io.Float.Input("structural_repulsion_boost", default=1.0, min=1.0, max=2.0, step=0.05, round=0.01, display_mode=io.NumberDisplay.slider, optional=True, display_name="structural_repulsion_boost"),
|
|
io.Int.Input("protect_zone_size", default=6, min=0, max=20, step=1, display_mode=io.NumberDisplay.slider, optional=True, display_name="protect_zone_size"),
|
|
io.Float.Input("time_smoothing_factor", default=0.2, min=0.0, max=1.0, step=0.05, round=0.01, display_mode=io.NumberDisplay.slider, optional=True, display_name="time_smoothing_factor"),
|
|
io.ClipVisionOutput.Input("clip_vision_output", optional=True),
|
|
],
|
|
outputs=[
|
|
io.Conditioning.Output(display_name="positive_high"),
|
|
io.Conditioning.Output(display_name="positive_low"),
|
|
io.Conditioning.Output(display_name="negative"),
|
|
io.Latent.Output(display_name="latent"),
|
|
],
|
|
)
|
|
|
|
@classmethod
|
|
def execute(cls, positive, negative, vae, width, height, length, batch_size, ref_images,
|
|
mode="NORMAL", ref_positions="", ref_strength_high=0.8, end_frame_strength_high=1.0,
|
|
ref_strength_low=0.2, start_frame_strength_low=1.0, end_frame_strength_low=1.0,
|
|
structural_repulsion_boost=1.0, protect_zone_size=6, time_smoothing_factor=0.2, clip_vision_output=None):
|
|
|
|
spacial_scale = vae.spacial_compression_encode()
|
|
latent_channels = vae.latent_channels
|
|
latent_t = ((length - 1) // 4) + 1
|
|
device = comfy.model_management.intermediate_device()
|
|
|
|
latent = torch.zeros([batch_size, latent_channels, latent_t,
|
|
height // spacial_scale, width // spacial_scale], device=device)
|
|
|
|
imgs = cls._resize_images(ref_images, width, height, device)
|
|
n_imgs = imgs.shape[0]
|
|
positions = cls._parse_positions(ref_positions, n_imgs, length)
|
|
|
|
def align_position(pos: int, total_frames: int) -> int:
|
|
latent_idx = pos // 4
|
|
aligned_pos = latent_idx * 4
|
|
aligned_pos = max(0, min(aligned_pos, total_frames - 1))
|
|
return aligned_pos
|
|
|
|
aligned_positions = [align_position(int(p), length) for p in positions]
|
|
|
|
for i in range(1, len(aligned_positions)):
|
|
if aligned_positions[i] <= aligned_positions[i-1] + 3:
|
|
aligned_positions[i] = min(aligned_positions[i-1] + 4, length - 1)
|
|
|
|
image = torch.ones((length, height, width, 3), device=device) * 0.5
|
|
mask_base = torch.ones((1, 1, latent_t * 4, latent.shape[-2], latent.shape[-1]), device=device)
|
|
|
|
mask_high_noise = mask_base.clone()
|
|
mask_low_noise = mask_base.clone()
|
|
|
|
for i, pos in enumerate(aligned_positions):
|
|
frame_idx = int(pos)
|
|
|
|
if i == 0:
|
|
image[frame_idx:frame_idx + 1] = imgs[i]
|
|
mask_high_noise[:, :, frame_idx:frame_idx + 4] = 0.0
|
|
|
|
low_start_mask_value = 1.0 - start_frame_strength_low
|
|
mask_low_noise[:, :, frame_idx:frame_idx + 4] = low_start_mask_value
|
|
elif i == n_imgs - 1:
|
|
image[-1:] = imgs[i]
|
|
|
|
mask_high_value = 1.0 - end_frame_strength_high
|
|
mask_high_noise[:, :, -4:] = mask_high_value
|
|
|
|
mask_low_value = 1.0 - end_frame_strength_low
|
|
mask_low_noise[:, :, -4:] = mask_low_value
|
|
else:
|
|
image[frame_idx:frame_idx + 1] = imgs[i]
|
|
|
|
mask_high_value = 1.0 - ref_strength_high
|
|
mask_high_noise[:, :, frame_idx:frame_idx + 4] = mask_high_value
|
|
|
|
concat_latent_image = vae.encode(image[:, :, :, :3])
|
|
|
|
concat_latent_image_high = concat_latent_image
|
|
|
|
if structural_repulsion_boost > 1.001 and length > 4 and n_imgs >= 2:
|
|
mask_h, mask_w = mask_high_noise.shape[-2], mask_high_noise.shape[-1]
|
|
boost_factor = structural_repulsion_boost - 1.0
|
|
|
|
def create_spatial_gradient(img1, img2):
|
|
if img1 is None or img2 is None:
|
|
return None
|
|
|
|
motion_diff = torch.abs(img2[0] - img1[0]).mean(dim=-1, keepdim=False)
|
|
motion_diff_4d = motion_diff.unsqueeze(0).unsqueeze(0)
|
|
motion_diff_scaled = F.interpolate(
|
|
motion_diff_4d,
|
|
size=(mask_h, mask_w),
|
|
mode='bilinear',
|
|
align_corners=False
|
|
)
|
|
|
|
motion_normalized = (motion_diff_scaled - motion_diff_scaled.min()) / (motion_diff_scaled.max() - motion_diff_scaled.min() + 1e-8)
|
|
|
|
spatial_gradient = 1.0 - motion_normalized * boost_factor * 2.5
|
|
spatial_gradient = torch.clamp(spatial_gradient, 0.02, 1.0)
|
|
return spatial_gradient[0, 0]
|
|
|
|
for i in range(n_imgs - 1):
|
|
pos1 = int(aligned_positions[i])
|
|
pos2 = int(aligned_positions[i + 1])
|
|
|
|
if pos2 > pos1 + 4:
|
|
start_end = pos1 + 4
|
|
ref_frame_start = pos2
|
|
transition_end = pos2
|
|
|
|
img1 = imgs[i:i+1].to(device)
|
|
img2 = imgs[i+1:i+2].to(device)
|
|
|
|
spatial_gradient = create_spatial_gradient(img1, img2)
|
|
|
|
if spatial_gradient is not None and transition_end > start_end:
|
|
transition_length = transition_end - start_end
|
|
|
|
for frame_idx in range(start_end, transition_end):
|
|
if frame_idx < ref_frame_start:
|
|
current_mask = mask_high_noise[:, :, frame_idx, :, :]
|
|
|
|
distance_to_start = (frame_idx - start_end) / max(1.0, transition_length)
|
|
distance_to_ref = abs(frame_idx - ref_frame_start) / max(1.0, protect_zone_size)
|
|
|
|
time_weight = 1.0 - distance_to_start * time_smoothing_factor
|
|
protect_weight = max(0.3, 1.0 - distance_to_ref) if frame_idx < ref_frame_start else 1.0
|
|
|
|
combined_weight = time_weight * protect_weight
|
|
adjusted_gradient = 1.0 - (1.0 - spatial_gradient) * combined_weight
|
|
|
|
mask_high_noise[:, :, frame_idx, :, :] = current_mask * adjusted_gradient
|
|
|
|
image_low_only = torch.ones((length, height, width, 3), device=device) * 0.5
|
|
|
|
for i, pos in enumerate(aligned_positions):
|
|
frame_idx = int(pos)
|
|
|
|
if i == 0:
|
|
if start_frame_strength_low > 0.0:
|
|
image_low_only[frame_idx:frame_idx + 1] = imgs[i]
|
|
elif i == n_imgs - 1:
|
|
if end_frame_strength_low > 0.0:
|
|
image_low_only[-1:] = imgs[i]
|
|
else:
|
|
if ref_strength_low > 0.0:
|
|
image_low_only[frame_idx:frame_idx + 1] = imgs[i]
|
|
|
|
concat_latent_image_low = vae.encode(image_low_only[:, :, :, :3])
|
|
|
|
# Print low noise mask distribution for debugging
|
|
if n_imgs >= 2:
|
|
print("\n=== Low Noise Mask Distribution ===")
|
|
print(f"Total frames: {length}, Reference frames: {n_imgs}")
|
|
print(f"Parameters: start_strength_low={start_frame_strength_low}, ref_strength_low={ref_strength_low}, end_strength_low={end_frame_strength_low}")
|
|
print("\nMask values at reference frame positions:")
|
|
for i, pos in enumerate(aligned_positions):
|
|
frame_idx = int(pos)
|
|
mask_samples = []
|
|
sample_frames = [max(0, frame_idx - 2), frame_idx, min(length - 1, frame_idx + 2), min(length - 1, frame_idx + 4)]
|
|
for sf in sample_frames:
|
|
if 0 <= sf < length:
|
|
mask_val = mask_low_noise[:, :, sf, :, :].mean().item()
|
|
mask_samples.append((sf, mask_val))
|
|
|
|
frame_type = "首帧" if i == 0 else ("末帧" if i == n_imgs - 1 else f"参考帧{i+1}")
|
|
print(f" {frame_type} (位置 {frame_idx}):")
|
|
for sf, mv in mask_samples:
|
|
print(f" Frame {sf}: mask = {mv:.3f}")
|
|
|
|
# Check for abrupt mask changes
|
|
print("\n检查参考帧之间的mask变化:")
|
|
for i in range(n_imgs - 1):
|
|
pos1 = int(aligned_positions[i])
|
|
pos2 = int(aligned_positions[i + 1])
|
|
|
|
# Check mask at boundaries
|
|
if pos1 + 4 < length:
|
|
mask_before = mask_low_noise[:, :, pos1 + 4, :, :].mean().item()
|
|
else:
|
|
mask_before = mask_low_noise[:, :, length - 1, :, :].mean().item()
|
|
|
|
if pos2 - 1 >= 0:
|
|
mask_after = mask_low_noise[:, :, pos2 - 1, :, :].mean().item()
|
|
else:
|
|
mask_after = mask_low_noise[:, :, 0, :, :].mean().item()
|
|
|
|
mask_at_ref = mask_low_noise[:, :, pos2, :, :].mean().item()
|
|
|
|
change = abs(mask_after - mask_before)
|
|
print(f" 参考帧{i+1} -> 参考帧{i+2}: frame {pos1+4} mask={mask_before:.3f}, frame {pos2-1} mask={mask_after:.3f}, 参考帧位置 mask={mask_at_ref:.3f}, 变化={change:.3f}")
|
|
if change > 0.2:
|
|
print(f" ⚠️ 警告: mask变化较大 ({change:.3f})，可能导致闪烁！")
|
|
print("=" * 40 + "\n")
|
|
|
|
mask_high_reshaped = mask_high_noise.view(1, mask_high_noise.shape[2] // 4, 4, mask_high_noise.shape[3], mask_high_noise.shape[4]).transpose(1, 2)
|
|
mask_low_reshaped = mask_low_noise.view(1, mask_low_noise.shape[2] // 4, 4, mask_low_noise.shape[3], mask_low_noise.shape[4]).transpose(1, 2)
|
|
|
|
positive_high_noise = node_helpers.conditioning_set_values(positive, {
|
|
"concat_latent_image": concat_latent_image_high,
|
|
"concat_mask": mask_high_reshaped
|
|
})
|
|
|
|
positive_low_noise = node_helpers.conditioning_set_values(positive, {
|
|
"concat_latent_image": concat_latent_image_low,
|
|
"concat_mask": mask_low_reshaped
|
|
})
|
|
|
|
|
|
negative_out = node_helpers.conditioning_set_values(negative, {
|
|
"concat_latent_image": concat_latent_image_high,
|
|
"concat_mask": mask_high_reshaped
|
|
})
|
|
|
|
if clip_vision_output is not None:
|
|
positive_low_noise = node_helpers.conditioning_set_values(positive_low_noise,
|
|
{"clip_vision_output": clip_vision_output})
|
|
|
|
negative_out = node_helpers.conditioning_set_values(negative_out,
|
|
{"clip_vision_output": clip_vision_output})
|
|
|
|
return io.NodeOutput(positive_high_noise, positive_low_noise, negative_out, {"samples": latent})
|
|
|
|
@classmethod
|
|
def _resize_images(cls, images, width, height, device):
|
|
images = images.to(device)
|
|
x = images.movedim(-1, 1)
|
|
x = comfy.utils.common_upscale(x, width, height, "bilinear", "center")
|
|
x = x.movedim(1, -1)
|
|
|
|
if x.shape[-1] == 4:
|
|
x = x[..., :3]
|
|
|
|
return x
|
|
|
|
@classmethod
|
|
def _parse_positions(cls, pos_str, n_imgs, length):
|
|
positions = []
|
|
s = (pos_str or "").strip()
|
|
|
|
if s:
|
|
try:
|
|
if s.startswith("["):
|
|
positions = json.loads(s)
|
|
else:
|
|
positions = [float(x.strip()) for x in s.split(",") if x.strip()]
|
|
except Exception:
|
|
positions = []
|
|
|
|
if not positions:
|
|
if n_imgs <= 1:
|
|
positions = [0]
|
|
else:
|
|
positions = [i * (length - 1) / (n_imgs - 1) for i in range(n_imgs)]
|
|
|
|
converted_positions = []
|
|
for p in positions:
|
|
if 0 <= p < 2.0:
|
|
converted_positions.append(int(p * (length - 1)))
|
|
else:
|
|
converted_positions.append(int(p))
|
|
|
|
converted_positions = [max(0, min(length - 1, p)) for p in converted_positions]
|
|
|
|
if len(converted_positions) > n_imgs:
|
|
converted_positions = converted_positions[:n_imgs]
|
|
elif len(converted_positions) < n_imgs:
|
|
converted_positions.extend([converted_positions[-1]] * (n_imgs - len(converted_positions)))
|
|
|
|
return converted_positions
|