tutorials / custom_nodes /dolphin_nodes /wan_i2v_continuity.py

Upload folder using huggingface_hub

9c98083 verified about 2 months ago

5.71 kB

import torch
import comfy.model_management
import comfy.utils
import node_helpers


class WanI2VContinuity_Dolphin:
    """
    Wan2.2 I2V 노드 — 클립 이어붙이기(연속성)용.
    기존 접근(WanImageToVideoSVIPro)의 문제: anchor 이미지를 따로 인코딩한 latent와,
    이전 클립 latent의 뒷부분을 그냥 잘라낸 조각(motion_latent)을 이어붙였다. WAN VAE는
    인과적(causal) 시간 압축이라 각 latent 프레임이 "그 시점까지의 시퀀스 문맥"을 전제로
    인코딩되는데, 서로 다른 문맥에서 나온 latent 조각을 이어붙이면 통계적으로 불연속이라
    화질이 튀고 동작 연속성도 못 살렸다.
    이 노드는 ComfyUI 코어 `WanImageToVideo`의 검증된 방식을 따른다:
    이전 클립에서 디코딩한 마지막 N프레임(픽셀)을 start_image로 받아, 회색(0.5) filler와
    함께 하나의 시퀀스로 만들어 **한 번의 통짜 VAE 인코딩**을 한다. 그래서 known 구간과
    생성 구간의 latent 문맥이 자연스럽게 이어진다. 마스크도 프레임수→latent수 변환을
    정확히 해서 known 프레임 전체를 고정한다.
    여기에 PainterI2V의 motion_amplitude(슬로우모션 방지)를 이식했다. 단 known 프레임은
    건드리지 않고, 뒤에 붙는 filler(생성 대상) latent만 마지막 known latent 대비 차이를
    증폭한다. motion_amplitude=1.0이면 코어 WanImageToVideo와 동일하게 동작한다.
    """

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "positive": ("CONDITIONING",),
                "negative": ("CONDITIONING",),
                "vae": ("VAE",),
                "width": ("INT", {"default": 832, "min": 16, "max": 4096, "step": 16}),
                "height": ("INT", {"default": 480, "min": 16, "max": 4096, "step": 16}),
                "length": ("INT", {"default": 81, "min": 1, "max": 9999, "step": 4}),
                "batch_size": ("INT", {"default": 1, "min": 1, "max": 4096}),
                "motion_amplitude": ("FLOAT", {"default": 1.0, "min": 1.0, "max": 2.0, "step": 0.05,
                                                 "tooltip": "1.0=코어 WanImageToVideo와 동일. >1.0이면 filler 구간에 슬로우모션 방지 증폭."}),
            },
            "optional": {
                "clip_vision_output": ("CLIP_VISION_OUTPUT",),
                "start_image": ("IMAGE", {"tooltip": "known 프레임(들). 이전 클립의 디코딩된 마지막 N프레임을 넣으면 연속성이 생김. 1장이면 일반 I2V."}),
            },
        }

    RETURN_TYPES = ("CONDITIONING", "CONDITIONING", "LATENT")
    RETURN_NAMES = ("positive", "negative", "latent")
    FUNCTION = "encode"
    CATEGORY = "Dolphin Node/Wan Video"

    def encode(self, positive, negative, vae, width, height, length, batch_size,
               motion_amplitude=1.0, clip_vision_output=None, start_image=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)
            n_known_frames = start_image.shape[0]

            # known 프레임 뒤를 회색(0.5) filler로 채운 뒤 통짜로 인코딩 (코어 방식)
            image = torch.ones(
                (length, height, width, start_image.shape[-1]),
                device=start_image.device, dtype=start_image.dtype,
            ) * 0.5
            image[:n_known_frames] = start_image

            concat_latent_image = vae.encode(image[:, :, :, :3])

            # 프레임수 -> latent수 (WAN VAE 4x 시간 압축)
            n_known_latents = ((n_known_frames - 1) // 4) + 1

            # motion_amplitude: filler(생성 대상) latent만 증폭. known 프레임은 그대로.
            if motion_amplitude > 1.0 and concat_latent_image.shape[2] > n_known_latents:
                base_latent = concat_latent_image[:, :, n_known_latents - 1:n_known_latents]
                filler = concat_latent_image[:, :, n_known_latents:]
                diff = filler - base_latent
                diff_mean = diff.mean(dim=(1, 3, 4), keepdim=True)
                diff_centered = diff - diff_mean
                scaled = base_latent + diff_centered * motion_amplitude + diff_mean
                scaled = torch.clamp(scaled, -6, 6)
                concat_latent_image = torch.cat(
                    [concat_latent_image[:, :, :n_known_latents], scaled], dim=2
                )

            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[:, :, :n_known_latents] = 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})

        return (positive, negative, {"samples": latent})

	import torch
	import comfy.model_management
	import comfy.utils
	import node_helpers


	class WanI2VContinuity_Dolphin:
	"""
	Wan2.2 I2V 노드 — 클립 이어붙이기(연속성)용.

	기존 접근(WanImageToVideoSVIPro)의 문제: anchor 이미지를 따로 인코딩한 latent와,
	이전 클립 latent의 뒷부분을 그냥 잘라낸 조각(motion_latent)을 이어붙였다. WAN VAE는
	인과적(causal) 시간 압축이라 각 latent 프레임이 "그 시점까지의 시퀀스 문맥"을 전제로
	인코딩되는데, 서로 다른 문맥에서 나온 latent 조각을 이어붙이면 통계적으로 불연속이라
	화질이 튀고 동작 연속성도 못 살렸다.

	이 노드는 ComfyUI 코어 `WanImageToVideo`의 검증된 방식을 따른다:
	이전 클립에서 디코딩한 마지막 N프레임(픽셀)을 start_image로 받아, 회색(0.5) filler와
	함께 하나의 시퀀스로 만들어 한 번의 통짜 VAE 인코딩을 한다. 그래서 known 구간과
	생성 구간의 latent 문맥이 자연스럽게 이어진다. 마스크도 프레임수→latent수 변환을
	정확히 해서 known 프레임 전체를 고정한다.

	여기에 PainterI2V의 motion_amplitude(슬로우모션 방지)를 이식했다. 단 known 프레임은
	건드리지 않고, 뒤에 붙는 filler(생성 대상) latent만 마지막 known latent 대비 차이를
	증폭한다. motion_amplitude=1.0이면 코어 WanImageToVideo와 동일하게 동작한다.
	"""

	@classmethod
	def INPUT_TYPES(cls):
	return {
	"required": {
	"positive": ("CONDITIONING",),
	"negative": ("CONDITIONING",),
	"vae": ("VAE",),
	"width": ("INT", {"default": 832, "min": 16, "max": 4096, "step": 16}),
	"height": ("INT", {"default": 480, "min": 16, "max": 4096, "step": 16}),
	"length": ("INT", {"default": 81, "min": 1, "max": 9999, "step": 4}),
	"batch_size": ("INT", {"default": 1, "min": 1, "max": 4096}),
	"motion_amplitude": ("FLOAT", {"default": 1.0, "min": 1.0, "max": 2.0, "step": 0.05,
	"tooltip": "1.0=코어 WanImageToVideo와 동일. >1.0이면 filler 구간에 슬로우모션 방지 증폭."}),
	},
	"optional": {
	"clip_vision_output": ("CLIP_VISION_OUTPUT",),
	"start_image": ("IMAGE", {"tooltip": "known 프레임(들). 이전 클립의 디코딩된 마지막 N프레임을 넣으면 연속성이 생김. 1장이면 일반 I2V."}),
	},
	}

	RETURN_TYPES = ("CONDITIONING", "CONDITIONING", "LATENT")
	RETURN_NAMES = ("positive", "negative", "latent")
	FUNCTION = "encode"
	CATEGORY = "Dolphin Node/Wan Video"

	def encode(self, positive, negative, vae, width, height, length, batch_size,
	motion_amplitude=1.0, clip_vision_output=None, start_image=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)
	n_known_frames = start_image.shape[0]

	# known 프레임 뒤를 회색(0.5) filler로 채운 뒤 통짜로 인코딩 (코어 방식)
	image = torch.ones(
	(length, height, width, start_image.shape[-1]),
	device=start_image.device, dtype=start_image.dtype,
	) * 0.5
	image[:n_known_frames] = start_image

	concat_latent_image = vae.encode(image[:, :, :, :3])

	# 프레임수 -> latent수 (WAN VAE 4x 시간 압축)
	n_known_latents = ((n_known_frames - 1) // 4) + 1

	# motion_amplitude: filler(생성 대상) latent만 증폭. known 프레임은 그대로.
	if motion_amplitude > 1.0 and concat_latent_image.shape[2] > n_known_latents:
	base_latent = concat_latent_image[:, :, n_known_latents - 1:n_known_latents]
	filler = concat_latent_image[:, :, n_known_latents:]
	diff = filler - base_latent
	diff_mean = diff.mean(dim=(1, 3, 4), keepdim=True)
	diff_centered = diff - diff_mean
	scaled = base_latent + diff_centered * motion_amplitude + diff_mean
	scaled = torch.clamp(scaled, -6, 6)
	concat_latent_image = torch.cat(
	[concat_latent_image[:, :, :n_known_latents], scaled], dim=2
	)

	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[:, :, :n_known_latents] = 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})

	return (positive, negative, {"samples": latent})