329 lines
15 KiB
Python
329 lines
15 KiB
Python
import torch
|
|
import comfy.clip_vision
|
|
import comfy.model_management
|
|
import comfy.utils
|
|
|
|
|
|
class SmartPainterLongVideo:
|
|
"""
|
|
Smart variation of PainterLongVideo that adds an optional CLIP Vision reference
|
|
image (with independent strength) on top of the existing pipeline.
|
|
|
|
Execute logic is identical to PainterLongVideo. The only addition is an
|
|
extra clip_vision_reference_image slot whose CLIP tokens are scaled by
|
|
cv_reference_strength and merged with the existing clip_vision_output on
|
|
the token axis before being injected into conditioning.
|
|
|
|
When clip_vision_reference_image is unconnected the node behaves exactly
|
|
like the original PainterLongVideo.
|
|
"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"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": 1000, "step": 1}),
|
|
"batch_size": ("INT", {"default": 1, "min": 1, "max": 4096}),
|
|
"motion_frames": ("INT", {"default": 5, "min": 1, "max": 20}),
|
|
"motion_amplitude": ("FLOAT", {"default": 1.15, "min": 1.0, "max": 2.0, "step": 0.05}),
|
|
},
|
|
"optional": {
|
|
"previous_video": ("IMAGE",),
|
|
"initial_reference_image": ("IMAGE", {
|
|
"tooltip": (
|
|
"Optional VAE-encoded reference appended to the conditioning's "
|
|
"reference_latents list (zero latent on the negative side).\n"
|
|
"\n"
|
|
"NOTE: Only effective on WAN checkpoints that ship with a ref_conv layer "
|
|
"(Fun Control, Animate, SCAIL, HuMo, VACE-merged). "
|
|
"Silently ignored by stock WAN 2.2 FLF2V / I2V / T2V and most custom merges. "
|
|
"For a reference-image channel that works on any base WAN 2.2 checkpoint, "
|
|
"use clip_vision_reference_image below instead.\n"
|
|
"\n"
|
|
"Leave unconnected to keep the original pipeline unchanged."
|
|
),
|
|
}),
|
|
"clip_vision_output": ("CLIP_VISION_OUTPUT",),
|
|
"clip_vision_reference_image": ("CLIP_VISION_OUTPUT", {
|
|
"tooltip": (
|
|
"Optional extra CLIP Vision slot for an off-screen reference image "
|
|
"(subject identity, clothing, object that isn't in start/end frames).\n"
|
|
"\n"
|
|
"Merged with clip_vision_output by concatenating their token sequences "
|
|
"along the token axis. Feeds the model's img_emb cross-attention pathway — "
|
|
"the only non-keyframe image-guidance channel available on stock WAN 2.2 "
|
|
"base checkpoints (no VACE / no ref_conv required).\n"
|
|
"\n"
|
|
"Provides semantic / style / identity guidance. NOT pixel-faithful — "
|
|
"good for 'keep the jacket this color', weak for 'reproduce this exact logo'.\n"
|
|
"\n"
|
|
"Leave unconnected to keep the original pipeline unchanged."
|
|
),
|
|
}),
|
|
"cv_reference_strength": ("FLOAT", {
|
|
"default": 1.0, "min": 0.0, "max": 3.0, "step": 0.05,
|
|
"tooltip": (
|
|
"Only active when clip_vision_reference_image is connected.\n"
|
|
"\n"
|
|
"Scales the reference image's CLIP token magnitudes before they are merged "
|
|
"with clip_vision_output. Higher values make the reference dominate the "
|
|
"cross-attention, lower values let the existing CLIP guidance win.\n"
|
|
"\n"
|
|
"0.0 = reference disabled (same as leaving clip_vision_reference_image "
|
|
"unconnected).\n"
|
|
"1.0 = equal weight with clip_vision_output tokens.\n"
|
|
"2.0-3.0 = reference boosted (useful if the reference subject is being "
|
|
"ignored by the model).\n"
|
|
"\n"
|
|
"Defaults to 1.0."
|
|
),
|
|
}),
|
|
"start_image": ("IMAGE",),
|
|
"end_image": ("IMAGE",),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("CONDITIONING", "CONDITIONING", "LATENT")
|
|
RETURN_NAMES = ("positive", "negative", "latent")
|
|
FUNCTION = "execute"
|
|
CATEGORY = "SmartHelperNodes/Wan"
|
|
DISPLAY_NAME = "Smart Painter Long Video"
|
|
DESCRIPTION = (
|
|
"Smart Painter Long Video: long-form WAN video conditioning with optional CLIP Vision "
|
|
"reference image and independent strength.\n"
|
|
"Identical execute logic to PainterLongVideo, plus an extra clip_vision_reference_image "
|
|
"slot for off-screen subject/style guidance through the model's img_emb cross-attention."
|
|
)
|
|
|
|
def execute(
|
|
self,
|
|
positive,
|
|
negative,
|
|
vae,
|
|
width,
|
|
height,
|
|
length,
|
|
batch_size,
|
|
motion_frames,
|
|
motion_amplitude=1.15,
|
|
previous_video=None,
|
|
initial_reference_image=None,
|
|
clip_vision_output=None,
|
|
clip_vision_reference_image=None,
|
|
cv_reference_strength=1.0,
|
|
start_image=None,
|
|
end_image=None,
|
|
):
|
|
device = comfy.model_management.intermediate_device()
|
|
latent_timesteps = ((length - 1) // 4) + 1
|
|
latent = torch.zeros([batch_size, 16, latent_timesteps, height // 8, width // 8], device=device)
|
|
|
|
# ---------- 1. 收集开关 ----------
|
|
has_prev = previous_video is not None
|
|
has_start = start_image is not None
|
|
has_end = end_image is not None
|
|
|
|
# ---------- 2. 无输入直接报错 ----------
|
|
if not has_prev and not has_start and not has_end:
|
|
raise RuntimeError("SmartPainterLongVideo: 必须接入 previous_video、start_image、end_image 中的至少一个！")
|
|
|
|
# ---------- 3. 构造统一的 image / mask ----------
|
|
# 默认值：灰色画布
|
|
image = torch.full((length, height, width, 3), 0.5, device=device, dtype=torch.float32)
|
|
mask = torch.ones((1, 1, latent_timesteps * 4, height // 8, width // 8), device=device, dtype=torch.float32)
|
|
|
|
# ---------- 4. 按优先级分支 ----------
|
|
if has_start or has_end:
|
|
# 分支 A/B：只要接了 start 或 end，就以「首尾帧」逻辑为准
|
|
if has_start:
|
|
start_img = comfy.utils.common_upscale(
|
|
start_image[:length].movedim(-1, 1), width, height, "bilinear", "center"
|
|
).movedim(1, -1)
|
|
actual_start_len = min(start_img.shape[0], length)
|
|
image[:actual_start_len] = start_img[:actual_start_len]
|
|
mask[:, :, :actual_start_len + 3] = 0.0 # 保护首帧（含缓冲）
|
|
else:
|
|
# 没给 start，但给了 end → 把 previous_video 最后一帧当首帧
|
|
if has_prev:
|
|
last_frame = previous_video[-1:].clone()
|
|
last_frame_resized = comfy.utils.common_upscale(
|
|
last_frame.movedim(-1, 1), width, height, "bilinear", "center"
|
|
).movedim(1, -1)
|
|
image[0] = last_frame_resized[0]
|
|
mask[:, :, :1 + 3] = 0.0
|
|
# 如果既没有 start 也没有 prev，却走到这里，说明只接了 end，也允许继续
|
|
|
|
if has_end:
|
|
end_img = comfy.utils.common_upscale(
|
|
end_image[-length:].movedim(-1, 1), width, height, "bilinear", "center"
|
|
).movedim(1, -1)
|
|
actual_end_len = min(end_img.shape[0], length)
|
|
image[-actual_end_len:] = end_img[-actual_end_len:]
|
|
mask[:, :, -actual_end_len:] = 0.0
|
|
|
|
concat_latent_image = vae.encode(image[:, :, :, :3])
|
|
mask = mask.view(1, mask.shape[2] // 4, 4, mask.shape[3], mask.shape[4]).transpose(1, 2)
|
|
|
|
# 运动参考：优先用 previous_video（如果给了）
|
|
ref_motion_latent = None
|
|
if has_prev and previous_video.shape[0] >= 2:
|
|
ref_motion = previous_video[-min(73, previous_video.shape[0]):].clone()
|
|
ref_motion = comfy.utils.common_upscale(
|
|
ref_motion.movedim(-1, 1), width, height, "bilinear", "center"
|
|
).movedim(1, -1)
|
|
if ref_motion.shape[0] < 73:
|
|
gray_fill = torch.ones([73, height, width, 3], device=device, dtype=ref_motion.dtype) * 0.5
|
|
gray_fill[-ref_motion.shape[0]:] = ref_motion
|
|
ref_motion = gray_fill
|
|
ref_motion_latent_temp = vae.encode(ref_motion[:, :, :, :3])
|
|
ref_motion_latent = ref_motion_latent_temp[:, :, -19:]
|
|
|
|
else:
|
|
# 分支 C：只接了 previous_video，走原逻辑
|
|
last_frame = previous_video[-1:].clone()
|
|
last_frame_resized = comfy.utils.common_upscale(
|
|
last_frame.movedim(-1, 1), width, height, "bilinear", "center"
|
|
).movedim(1, -1)
|
|
|
|
image_seq = torch.full(
|
|
(length, height, width, last_frame_resized.shape[-1]),
|
|
0.5, device=last_frame_resized.device, dtype=last_frame_resized.dtype
|
|
)
|
|
image_seq[0] = last_frame_resized[0]
|
|
concat_latent_image = vae.encode(image_seq[:, :, :, :3])
|
|
|
|
mask = torch.ones((1, 1, latent_timesteps, height // 8, width // 8),
|
|
device=device, dtype=last_frame_resized.dtype)
|
|
mask[:, :, 0] = 0.0
|
|
|
|
# 运动幅度增强
|
|
if motion_amplitude > 1.0:
|
|
base_latent = concat_latent_image[:, :, 0:1]
|
|
gray_latent = concat_latent_image[:, :, 1:]
|
|
diff = gray_latent - base_latent
|
|
diff_mean = diff.mean(dim=(1, 3, 4), keepdim=True)
|
|
diff_centered = diff - diff_mean
|
|
scaled_latent = base_latent + diff_centered * motion_amplitude + diff_mean
|
|
scaled_latent = torch.clamp(scaled_latent, -6, 6)
|
|
concat_latent_image = torch.cat([base_latent, scaled_latent], dim=2)
|
|
|
|
# 运动参考帧
|
|
ref_motion = previous_video[-motion_frames:].clone()
|
|
if ref_motion.shape[0] > 73:
|
|
ref_motion = ref_motion[-73:]
|
|
ref_motion_resized = comfy.utils.common_upscale(
|
|
ref_motion.movedim(-1, 1), width, height, "bilinear", "center"
|
|
).movedim(1, -1)
|
|
if ref_motion_resized.shape[0] < 73:
|
|
gray_fill = torch.ones([73, height, width, 3], device=ref_motion_resized.device,
|
|
dtype=ref_motion_resized.dtype) * 0.5
|
|
gray_fill[-ref_motion_resized.shape[0]:] = ref_motion_resized
|
|
ref_motion_resized = gray_fill
|
|
ref_motion_latent_temp = vae.encode(ref_motion_resized[:, :, :, :3])
|
|
ref_motion_latent = ref_motion_latent_temp[:, :, -19:]
|
|
|
|
# ---------- 5. 构造 reference_latents ----------
|
|
ref_latents = []
|
|
# (a) 上一段最后一帧（始终添加）
|
|
last_frame_for_ref = previous_video[-1:] if has_prev else image[0:1]
|
|
last_frame_for_ref = comfy.utils.common_upscale(
|
|
last_frame_for_ref.movedim(-1, 1), width, height, "bilinear", "center"
|
|
).movedim(1, -1)
|
|
last_latent = vae.encode(last_frame_for_ref[:, :, :, :3])
|
|
ref_latents.append(last_latent)
|
|
|
|
# (b) 初始参考图（如有）
|
|
if initial_reference_image is not None:
|
|
init_img = initial_reference_image[:1]
|
|
init_img_resized = comfy.utils.common_upscale(
|
|
init_img.movedim(-1, 1), width, height, "bilinear", "center"
|
|
).movedim(1, -1)
|
|
init_latent = vae.encode(init_img_resized[:, :, :, :3])
|
|
ref_latents.append(init_latent)
|
|
|
|
# ---------- 6. conditioning 注入 ----------
|
|
def inject_conditioning(cond, values_dict):
|
|
new_cond = []
|
|
for c_tensor, c_dict in cond:
|
|
new_dict = c_dict.copy()
|
|
new_dict.update(values_dict)
|
|
new_cond.append([c_tensor, new_dict])
|
|
return new_cond
|
|
|
|
def append_conditioning(cond, key, value_list):
|
|
new_cond = []
|
|
for c_tensor, c_dict in cond:
|
|
new_dict = c_dict.copy()
|
|
if key in new_dict:
|
|
new_dict[key] = new_dict[key] + value_list
|
|
else:
|
|
new_dict[key] = value_list
|
|
new_cond.append([c_tensor, new_dict])
|
|
return new_cond
|
|
|
|
shared_values = {
|
|
"concat_latent_image": concat_latent_image,
|
|
"concat_mask": mask,
|
|
}
|
|
if ref_motion_latent is not None:
|
|
shared_values["reference_motion"] = ref_motion_latent
|
|
|
|
pos_out = inject_conditioning(positive, shared_values)
|
|
neg_out = inject_conditioning(negative, shared_values)
|
|
|
|
pos_out = append_conditioning(pos_out, "reference_latents", ref_latents)
|
|
neg_ref_latents = [torch.zeros_like(r) for r in ref_latents]
|
|
neg_out = append_conditioning(neg_out, "reference_latents", neg_ref_latents)
|
|
|
|
scaled_reference = self._scale_clip_vision(
|
|
clip_vision_reference_image, cv_reference_strength,
|
|
)
|
|
merged_clip_vision = self._merge_clip_vision(clip_vision_output, scaled_reference)
|
|
if merged_clip_vision is not None:
|
|
pos_out = inject_conditioning(pos_out, {"clip_vision_output": merged_clip_vision})
|
|
neg_out = inject_conditioning(neg_out, {"clip_vision_output": merged_clip_vision})
|
|
|
|
latent_out = {"samples": latent}
|
|
return (pos_out, neg_out, latent_out)
|
|
|
|
# ------------------------------------------------------------------
|
|
@staticmethod
|
|
def _merge_clip_vision(*outputs):
|
|
valid = [o for o in outputs if o is not None]
|
|
if not valid:
|
|
return None
|
|
if len(valid) == 1:
|
|
return valid[0]
|
|
states = torch.cat([o.penultimate_hidden_states for o in valid], dim=-2)
|
|
result = comfy.clip_vision.Output()
|
|
result.penultimate_hidden_states = states
|
|
return result
|
|
|
|
# ------------------------------------------------------------------
|
|
@staticmethod
|
|
def _scale_clip_vision(output, strength):
|
|
if output is None or strength <= 0.0:
|
|
return None
|
|
if abs(strength - 1.0) < 1e-6:
|
|
return output
|
|
scaled = comfy.clip_vision.Output()
|
|
scaled.penultimate_hidden_states = output.penultimate_hidden_states * strength
|
|
for attr in ("image_embeds", "last_hidden_state", "mm_projected", "image_sizes"):
|
|
if hasattr(output, attr):
|
|
setattr(scaled, attr, getattr(output, attr))
|
|
return scaled
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"SmartPainterLongVideo": SmartPainterLongVideo,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"SmartPainterLongVideo": "Smart Painter Long Video",
|
|
}
|