ComfyUI node to split Mochi latent frames

mochi_split_latent.py

# By https://github.com/blepping/

import torch

class SplitMochiVideoLatent:
    DESCRIPTION = "Hack for Mochi latents to split the frames. Mochi has 6x temporal compression so a latent frame is worth 6 frames in your video. Splitting up a Mochi latent may reduce quality."
    FUNCTION = "go"
    CATEGORY = "latent/mochi"
    RETURN_TYPES = ("LATENT",)

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "latent": ("LATENT",),
                "latent_frame_index": ("INT", {"min": 0}),
                "latent_frame_count": ("INT", {"min": 0}),
            }
        }

    @classmethod
    def go(cls, *, latent_frame_index, latent_frame_count, latent):
        samples = latent["samples"]
        shape = samples.shape
        if samples.ndim != 5 or shape[1] != 12:
            raise ValueError("Bad latent shape, may not be a Mochi latent")
        frames = shape[2]
        if latent_frame_index >= frames:
            raise ValueError("Latent frame index is greater than the number of latent frames. Note that Mochi has 6x temporal compression.")
        result = samples[:, :, latent_frame_index:latent_frame_index + latent_frame_count, :, :].detach().clone()
        return ({"samples": result},)



NODE_CLASS_MAPPINGS = {
    "SplitMochiVideoLatent": SplitMochiVideoLatent,
}