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