143 lines
7.3 KiB
Python
143 lines
7.3 KiB
Python
import torch
|
|
import comfy.model_management
|
|
import comfy.utils
|
|
import node_helpers
|
|
import comfy.clip_vision
|
|
|
|
# ==========================================
|
|
# 核心辅助函数：VAE 空间分块编码
|
|
# ==========================================
|
|
def _encode_vae(vae, pixels, tile_size):
|
|
if tile_size == 0:
|
|
return vae.encode(pixels[:, :, :, :3])
|
|
else:
|
|
# VAE 编码分块，降低图片转 Latent 瞬间的显存峰值
|
|
return vae.encode_tiled(pixels[:, :, :, :3], tile_x=tile_size, tile_y=tile_size, overlap=32, tile_t=256, overlap_t=8)
|
|
|
|
|
|
# ==========================================
|
|
# 节点 1：WAN 图像到视频 (带 VAE 分块)
|
|
# ==========================================
|
|
class XB_WanImageToVideo:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"positive": ("CONDITIONING",),
|
|
"negative": ("CONDITIONING",),
|
|
"vae": ("VAE",),
|
|
"width": ("INT", {"default": 480, "min": 16, "max": 8192, "step": 16}),
|
|
"height": ("INT", {"default": 832, "min": 16, "max": 8192, "step": 16}),
|
|
"length": ("INT", {"default": 81, "min": 1, "max": 9999, "step": 4}),
|
|
"batch_size": ("INT", {"default": 1, "min": 1, "max": 4096}),
|
|
"vae_tile_size": ("INT", {"default": 320, "min": 0, "max": 4096, "step": 16, "tooltip": "VAE编码分块大小，0表示不分块(默认0)"}),
|
|
},
|
|
"optional": {
|
|
"clip_vision_output": ("CLIP_VISION_OUTPUT",),
|
|
"start_image": ("IMAGE",),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("CONDITIONING", "CONDITIONING", "LATENT")
|
|
RETURN_NAMES = ("positive", "negative", "latent")
|
|
FUNCTION = "process"
|
|
CATEGORY = "小白工具箱/WAN解码分块"
|
|
|
|
def process(self, positive, negative, vae, width, height, length, batch_size, vae_tile_size=0, start_image=None, clip_vision_output=None):
|
|
latent = torch.zeros([batch_size, 16, ((length - 1) // 4) + 1, height // 8, width // 8], device=comfy.model_management.intermediate_device())
|
|
|
|
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)
|
|
image = torch.ones((length, height, width, start_image.shape[-1]), device=start_image.device, dtype=start_image.dtype) * 0.5
|
|
image[:start_image.shape[0]] = start_image
|
|
|
|
concat_latent_image = _encode_vae(vae, image, vae_tile_size)
|
|
|
|
mask = torch.ones((1, 1, latent.shape[2], concat_latent_image.shape[-2], concat_latent_image.shape[-1]), device=start_image.device, dtype=start_image.dtype)
|
|
mask[:, :, :((start_image.shape[0] - 1) // 4) + 1] = 0.0
|
|
|
|
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})
|
|
|
|
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})
|
|
|
|
out_latent = {"samples": latent}
|
|
return (positive, negative, out_latent)
|
|
|
|
|
|
# ==========================================
|
|
# 节点 2：WAN 首尾帧视频 (带 VAE 分块)
|
|
# ==========================================
|
|
class XB_WanFirstLastFrameToVideo:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"positive": ("CONDITIONING",),
|
|
"negative": ("CONDITIONING",),
|
|
"vae": ("VAE",),
|
|
"width": ("INT", {"default": 480, "min": 16, "max": 8192, "step": 16}),
|
|
"height": ("INT", {"default": 832, "min": 16, "max": 8192, "step": 16}),
|
|
"length": ("INT", {"default": 81, "min": 1, "max": 9999, "step": 4}),
|
|
"batch_size": ("INT", {"default": 1, "min": 1, "max": 4096}),
|
|
"vae_tile_size": ("INT", {"default": 320, "min": 0, "max": 4096, "step": 16}),
|
|
},
|
|
"optional": {
|
|
"clip_vision_start_image": ("CLIP_VISION_OUTPUT",),
|
|
"clip_vision_end_image": ("CLIP_VISION_OUTPUT",),
|
|
"start_image": ("IMAGE",),
|
|
"end_image": ("IMAGE",),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("CONDITIONING", "CONDITIONING", "LATENT")
|
|
RETURN_NAMES = ("positive", "negative", "latent")
|
|
FUNCTION = "process"
|
|
CATEGORY = "小白工具箱/WAN解码分块"
|
|
|
|
def process(self, positive, negative, vae, width, height, length, batch_size, vae_tile_size=0, start_image=None, end_image=None, clip_vision_start_image=None, clip_vision_end_image=None):
|
|
spacial_scale = vae.spacial_compression_encode()
|
|
latent = torch.zeros([batch_size, vae.latent_channels, ((length - 1) // 4) + 1, height // spacial_scale, width // spacial_scale], device=comfy.model_management.intermediate_device())
|
|
|
|
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)
|
|
|
|
image = torch.ones((length, height, width, 3)) * 0.5
|
|
mask = torch.ones((1, 1, latent.shape[2] * 4, latent.shape[-2], latent.shape[-1]))
|
|
|
|
if start_image is not None:
|
|
image[:start_image.shape[0]] = start_image
|
|
mask[:, :, :start_image.shape[0] + 3] = 0.0
|
|
|
|
if end_image is not None:
|
|
image[-end_image.shape[0]:] = end_image
|
|
mask[:, :, -end_image.shape[0]:] = 0.0
|
|
|
|
concat_latent_image = _encode_vae(vae, image, vae_tile_size)
|
|
|
|
mask = mask.view(1, mask.shape[2] // 4, 4, mask.shape[3], mask.shape[4]).transpose(1, 2)
|
|
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_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})
|
|
|
|
out_latent = {"samples": latent}
|
|
return (positive, negative, out_latent) |