WAN 2.2: looping video mode?

[noisefloordev]

I was experimenting with WanFirstLastFrameToVideo with the WAN2.2 i2v model. It works really well with it.

I tried an experiment: instead of providing a start and end frame, provide a start frame and a middle frame, then dupe the start frame at the end to create loopable videos. I didn't expect it to work at all, but it almost does! The result is teasingly close: it does hit the middle frame and return to the start.

But, it has an artifact: the middle frame is washed out, with a couple frames around it fading the difference. I'm guessing some of the 50% initial grey is leaking in. I don't know if there's a way to fix this or not, since I don't actually understand how the masking works. This is just mimicking what start_frame does. I've thrown a bunch of darts, like clearing the mask at that frame instead of the 4-frame group like start_frame does (that makes it worse).

It's probably just "you can't do that", but it's so close I figured I'd share it in case anyone has a clue.

class WanFirstMiddleLastFrameToVideo:
    @classmethod
    def INPUT_TYPES(s):
        return {"required": {"positive": ("CONDITIONING", ),
                             "negative": ("CONDITIONING", ),
                             "vae": ("VAE", ),
                             "width": ("INT", {"default": 832, "min": 16, "max": nodes.MAX_RESOLUTION, "step": 16}),
                             "height": ("INT", {"default": 480, "min": 16, "max": nodes.MAX_RESOLUTION, "step": 16}),
                             "length": ("INT", {"default": 81, "min": 1, "max": nodes.MAX_RESOLUTION, "step": 4}),
                             "batch_size": ("INT", {"default": 1, "min": 1, "max": 4096}),
                },
                "optional": {"clip_vision_start_image": ("CLIP_VISION_OUTPUT", ),
                             "clip_vision_end_image": ("CLIP_VISION_OUTPUT", ),
                             "start_image": ("IMAGE", ),
                             "end_image": ("IMAGE", ),
                             # NEW: toggle to place end frames in the middle and loop start->end
                             "loop_mid": ("BOOLEAN", {"default": False}),
                }}

    RETURN_TYPES = ("CONDITIONING", "CONDITIONING", "LATENT")
    RETURN_NAMES = ("positive", "negative", "latent")
    FUNCTION = "encode"

    CATEGORY = "conditioning/video_models"

    def encode(self, positive, negative, vae, width, height, length, batch_size, start_image=None, end_image=None, clip_vision_start_image=None, clip_vision_end_image=None, loop_mid=False):
        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)
        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 not loop_mid:
            if end_image is not None:
                image[-end_image.shape[0]:] = end_image
                mask[:, :, -end_image.shape[0]:] = 0.0
        else:
            if end_image is not None:
                mid = length // 2
                mid &= ~3
                print(f'Mid frame: {mid}')
                image[mid:mid + end_image.shape[0]] = end_image
                mask[:, :, mid:mid + end_image.shape[0] + 3] = 0.0

            # Lock the last frame to the first
            if start_image is not None:
                image[-start_image.shape[0]:] = start_image
                mask[:, :, -start_image.shape[0]:] = 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)
        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 = {}
        out_latent["samples"] = latent
        return (positive, negative, out_latent)