init

keyframe/interface.py

@@ -0,0 +1,57 @@
+from dataclasses import dataclass, field
+from typing import Union
+
+import torch
+
+
+class KeyframePart:
+    def __init__(self, batch_index: int, image: torch.Tensor, denoise: float) -> None:
+        self.batch_index = batch_index
+        self.denoise = denoise
+        self.image = image
+
+
+class KeyframePartGroup:
+    def __init__(self) -> None:
+        self.keyframes: list[KeyframePart] = []
+
+    def add(self, keyframe: KeyframePart) -> None:
+        added = False
+        for i in range(len(self.keyframes)):
+            if self.keyframes[i].batch_index == keyframe.batch_index:
+                self.keyframes[i] = keyframe
+                added = True
+                break
+        if not added:
+            self.keyframes.append(keyframe)
+        self.keyframes.sort(key=lambda k: k.batch_index)
+
+    def get_index(self, index: int) -> Union[KeyframePart, None]:
+        try:
+            return self.keyframes[index]
+        except IndexError:
+            return None
+
+    def __getitem__(self, index) -> KeyframePart:
+        return self.keyframes[index]
+
+    def is_empty(self) -> bool:
+        return len(self.keyframes) == 0
+
+
+@dataclass
+class ModelInjectParam:
+    keyframe_part_group: KeyframePartGroup
+    latent: dict = field(default_factory=dict, repr=False)
+    seed: int = 0
+    steps: int = 0
+    scheduler: str = 'normal'
+    denoise: float = 0
+    noise: torch.Tensor = field(default=None, repr=False)
+
+    def reset(self):
+        self.seed: int = 0
+        self.steps: int = 0
+        self.scheduler: str = 'normal'
+        self.denoise: float = 0
+        self.noise: torch.Tensor = None

keyframe/nodes.py

@@ -0,0 +1,166 @@
+import numpy as np
+
+from .interface import KeyframePartGroup, KeyframePart, ModelInjectParam
+from .sampling import keyframe_sample_factory
+
+from .util import inject_model, print
+from .samples import inject_samples
+import comfy.sample as comfy_sample
+
+inject_samples()
+comfy_sample.sample = keyframe_sample_factory(comfy_sample.sample)
+
+
+class KeyframePartNode:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "image": ("IMAGE",),
+                "batch_index": ("INT", {"default": 0, "min": 0, "max": 9999, "step": 1}),
+                "denoise": ("FLOAT", {"default": 0.1, "min": 0.001, "max": 1.0, "step": 0.01}),
+            },
+            "optional": {
+                "part": ("LATENT_KEYFRAME_PART",),
+            }
+        }
+
+    RETURN_TYPES = ("LATENT_KEYFRAME_PART",)
+    RETURN_NAMES = ("part",)
+    FUNCTION = "load_keyframe_part"
+
+    CATEGORY = "KeyframePart"
+
+    def load_keyframe_part(self, image, batch_index, denoise, part=None):
+        if not part:
+            part = KeyframePartGroup()
+        keyframe = KeyframePart(batch_index, image, denoise)
+        part.add(keyframe)
+        return (part,)
+
+
+class KeyframeInterpolationPartNode:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "image": ("IMAGE",),
+                "batch_index_from": ("INT", {"default": 0, "min": 0, "max": 9999, "step": 1}),
+                "batch_index_to": ("INT", {"default": 4, "min": 1, "max": 9999, "step": 1}),
+                "denoise_from": ("FLOAT", {"default": 0.1, "min": 0.001, "max": 1.0, "step": 0.01}),
+                "denoise_to": ("FLOAT", {"default": 1.0, "min": 0.001, "max": 1.0, "step": 0.01}),
+                "interpolation": (["linear", "ease-in", "ease-out", "ease-in-out"],),
+            },
+            "optional": {
+                "part": ("LATENT_KEYFRAME_PART",),
+            }
+        }
+
+    RETURN_TYPES = ("LATENT_KEYFRAME_PART",)
+    RETURN_NAMES = ("part",)
+    FUNCTION = "load_keyframe_part"
+
+    CATEGORY = "KeyframeInterpolationPartNode"
+
+    def load_keyframe_part(self,
+                           image,
+                           batch_index_from,
+                           batch_index_to,
+                           denoise_from,
+                           denoise_to,
+                           interpolation,
+                           part=None):
+        if batch_index_from >= batch_index_to:
+            raise ValueError("batch_index_from must be less than batch_index_to")
+
+        if not part:
+            part = KeyframePartGroup()
+        current_group = KeyframePartGroup()
+
+        steps = batch_index_to - batch_index_from
+        diff = denoise_to - denoise_from
+        if interpolation == "linear":
+            weights = np.linspace(denoise_from, denoise_to, steps)
+        elif interpolation == "ease-in":
+            index = np.linspace(0, 1, steps)
+            weights = diff * np.power(index, 2) + denoise_from
+        elif interpolation == "ease-out":
+            index = np.linspace(0, 1, steps)
+            weights = diff * (1 - np.power(1 - index, 2)) + denoise_from
+        elif interpolation == "ease-in-out":
+            index = np.linspace(0, 1, steps)
+            weights = diff * ((1 - np.cos(index * np.pi)) / 2) + denoise_from
+
+        for i in range(steps):
+            keyframe = KeyframePart(batch_index_from + i, image, float(weights[i]))
+            current_group.add(keyframe)
+
+        # replace values with prev_latent_keyframes
+        for latent_keyframe in part.keyframes:
+            current_group.add(latent_keyframe)
+
+        return (current_group,)
+
+
+class KeyframeApplyNode:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "model": ("MODEL",),
+                "latent": ("LATENT",),
+                "part_group": ("LATENT_KEYFRAME_PART",),
+                "vae": ("VAE",),
+            }
+        }
+
+    RETURN_TYPES = ("MODEL", "LATENT",)
+    RETURN_NAMES = ("model", "latent",)
+    FUNCTION = "apply_latent_keyframe"
+
+    CATEGORY = "LatentKeyframeApply"
+
+    @staticmethod
+    def vae_encode_crop_pixels(pixels):
+        x = (pixels.shape[1] // 8) * 8
+        y = (pixels.shape[2] // 8) * 8
+        if pixels.shape[1] != x or pixels.shape[2] != y:
+            x_offset = (pixels.shape[1] % 8) // 2
+            y_offset = (pixels.shape[2] % 8) // 2
+            pixels = pixels[:, x_offset:x + x_offset, y_offset:y + y_offset, :]
+        return pixels
+
+    def encode(self, vae, pixels):
+        pixels = self.vae_encode_crop_pixels(pixels)
+        t = vae.encode(pixels[:, :, :, :3])
+        return t
+
+    def apply_latent_keyframe(self, model, latent, part_group: KeyframePartGroup, vae):
+        # 预处理latent,把图片替换进去
+        for part in part_group.keyframes:
+            latent['samples'][part.batch_index] = self.encode(vae, part.image)
+            print(f"apply keyframe {part.batch_index}:{part.denoise}")
+
+        # 注入参数,后续处理
+
+        inject_param = ModelInjectParam(part_group, latent)
+        inject_model(model.model, inject_param)
+
+        return (model, latent,)
+
+
+# NODE MAPPING
+NODE_CLASS_MAPPINGS = {
+    # Keyframes
+    "KeyframePart": KeyframePartNode,
+    "KeyframeInterpolationPart": KeyframeInterpolationPartNode,
+    "KeyframeApply": KeyframeApplyNode,
+
+}
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+    # Keyframes
+    "KeyframePart": "Keyframe Part",
+    "KeyframeInterpolationPart": "Keyframe Interpolation Part",
+    "KeyframeApply": "Keyframe Apply",
+}

keyframe/samples.py

@@ -0,0 +1,104 @@
+import torch
+from tqdm.auto import trange
+
+import comfy.samplers
+from comfy.k_diffusion import sampling as k_diffusion_sampling
+from comfy.k_diffusion.sampling import to_d, default_noise_sampler
+from .util import print, is_injected_model, get_injected_model, generate_sigmas, generate_noise, get_ancestral_step
+
+CUSTOM_SAMPLERS = [
+    'k_euler', 'k_euler_a'
+]
+
+
+def inject_samples():
+    comfy.samplers.SAMPLER_NAMES.extend(CUSTOM_SAMPLERS)
+    k_diffusion_sampling.sample_k_euler = sample_k_euler
+    k_diffusion_sampling.sample_k_euler_a = sample_k_euler_a
+    print(f'Injected samplers: {CUSTOM_SAMPLERS}')
+
+
+def get_sigmas_noise(model_wrap, x, noise, latent_image, sigmas, scheduler, steps, part_group):
+    sigmas = generate_sigmas(model_wrap.inner_model, x, sigmas, scheduler, steps, part_group, sigmas.device)
+    noise = noise.to(x.device)
+    latent_image = latent_image.to(x.device)
+    for i in range(noise.shape[0]):
+        noise[i] = generate_noise(model_wrap, sigmas[i], noise[i])
+
+    if latent_image is not None:
+        latent_image = model_wrap.inner_model.process_latent_in(latent_image)
+        noise += latent_image
+    return sigmas, noise
+
+
+@torch.no_grad()
+def sample_k_euler(model, x, sigmas, extra_args=None, callback=None, disable=None, s_churn=0., s_tmin=0.,
+                   s_tmax=float('inf'), s_noise=1.):
+    model_wrap = model.inner_model
+    real_model = model_wrap.inner_model
+
+    if not is_injected_model(real_model):
+        raise Exception("model is not injected,please use LatentKeyframeApply node to inject model")
+
+    inject_param = get_injected_model(real_model)
+    latent_image = inject_param.latent['samples']
+
+    sigmas, noise = get_sigmas_noise(model_wrap, x, inject_param.noise, latent_image, sigmas, inject_param.scheduler,
+                                     inject_param.steps, inject_param.keyframe_part_group)
+
+    extra_args = {} if extra_args is None else extra_args
+    s_tmin = s_tmin * x.new_ones([sigmas.shape[0]])
+    s_tmax = s_tmax * x.new_ones([sigmas.shape[0]])
+    gammas = s_tmax * x.new_ones([sigmas.shape[0]])
+    for i in trange(sigmas.shape[1] - 1, disable=disable):
+        for j, sigma in enumerate(sigmas):
+            gammas[j] = min(s_churn / (sigmas.shape[0] - 1), 2 ** 0.5 - 1) if s_tmin[j] <= sigma[i] <= s_tmax[j] else 0.
+        sigma_hat = sigmas[:, i] * (gammas + 1)
+
+        for j, gamma in enumerate(gammas):
+            if gamma > 0:
+                eps = torch.randn_like(noise) * s_noise
+                noise[j] = noise[j] + eps * (sigma_hat ** 2 - sigmas[i] ** 2) ** 0.5
+
+        denoised = model(noise, sigma_hat, **extra_args)
+        d = to_d(noise, sigma_hat, denoised)
+        if callback is not None:
+            callback({'x': noise, 'i': i, 'sigma': sigma_hat, 'denoised': denoised})
+        dt = sigmas[:, i + 1] - sigma_hat
+        noise += d * dt.view(d.shape[0], 1, 1, 1)
+    return noise
+
+
+@torch.no_grad()
+def sample_k_euler_a(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1.,
+                     noise_sampler=None):
+
+    model_wrap = model.inner_model
+    real_model = model_wrap.inner_model
+
+    if not is_injected_model(real_model):
+        raise Exception("model is not injected,please use LatentKeyframeApply node to inject model")
+
+    inject_param = get_injected_model(real_model)
+    latent_image = inject_param.latent['samples']
+
+    sigmas, noise = get_sigmas_noise(model_wrap, x, inject_param.noise, latent_image, sigmas, inject_param.scheduler,
+                                     inject_param.steps, inject_param.keyframe_part_group)
+
+    extra_args = {} if extra_args is None else extra_args
+    noise_sampler = default_noise_sampler(noise) if noise_sampler is None else noise_sampler
+
+    for i in trange(sigmas.shape[1] - 1, disable=disable):
+        s_in = sigmas[:, i]
+        denoised = model(noise, s_in, **extra_args)
+        sigma_down, sigma_up = get_ancestral_step(sigmas[:, i], sigmas[:, i + 1], eta=eta)
+        if callback is not None:
+            callback({'x': noise, 'i': i, 'sigma': s_in, 'denoised': denoised})
+        d = to_d(noise, s_in, denoised)
+        # Euler method
+        dt = sigma_down - s_in
+        noise += d * dt.view(d.shape[0], 1, 1, 1)
+        for j, sigma in enumerate(sigmas):
+            if sigma[i + 1] > 0:
+                noise[j] = noise[j] + noise_sampler(sigma[i], sigma[i + 1])[j] * s_noise * sigma_up[j]
+    return noise

keyframe/sampling.py

@@ -0,0 +1,24 @@
+from typing import Callable
+
+from comfy.model_patcher import ModelPatcher
+from .util import is_injected_model, print, get_injected_model
+
+
+def keyframe_sample_factory(orig_comfy_sample: Callable) -> Callable:
+    def keyframe_sample(model: ModelPatcher, *args, **kwargs):
+        if not is_injected_model(model.model):
+            return orig_comfy_sample(model, *args, **kwargs)
+        inject_param = get_injected_model(model.model)
+        try:
+            inject_param.reset()
+
+            inject_param.noise = args[0]
+            inject_param.steps = args[1]
+            inject_param.scheduler = args[4]
+            inject_param.denoise = kwargs.get('denoise', None)
+            inject_param.seed = kwargs.get('seed', None)
+            return orig_comfy_sample(model, *args, **kwargs)
+        finally:
+            inject_param.reset()
+
+    return keyframe_sample

keyframe/util.py

@@ -0,0 +1,72 @@
+import math
+
+import torch
+
+from comfy.samplers import calculate_sigmas_scheduler
+from rich.console import Console
+
+KEYFRAME_INJECTED_ATTR = "keyframe_injected"
+console = Console(color_system="truecolor", force_terminal=True)
+
+
+def inject_model(model, inject_param):
+    # 注入模型参数
+    setattr(model, KEYFRAME_INJECTED_ATTR, inject_param)
+    return model
+
+
+def is_injected_model(model):
+    return hasattr(model, KEYFRAME_INJECTED_ATTR)
+
+
+def get_injected_model(model):
+    return getattr(model, KEYFRAME_INJECTED_ATTR)
+
+
+def clear_injected_model(model):
+    if is_injected_model(model):
+        delattr(model, KEYFRAME_INJECTED_ATTR)
+
+
+def print(msg, *args, **kwargs):
+    msg = f'[bold red]Keyframe[/bold red] [green]{msg}[/green]'
+    console.print(msg)
+
+
+def max_denoise(model_wrap, sigmas):
+    max_sigma = float(model_wrap.inner_model.model_sampling.sigma_max)
+    sigma = float(sigmas[0])
+    return math.isclose(max_sigma, sigma, rel_tol=1e-05) or sigma > max_sigma
+
+
+def generate_sigmas(real_model, x, origin_sigmas, scheduler, steps, part_group, device):
+    batch_size = x.shape[0]
+    new_sigmas = origin_sigmas.unsqueeze(0).repeat(batch_size, 1)
+
+    for part in part_group:
+        if part.denoise is None or part.denoise > 0.9999:
+            new_sigmas[part.batch_index] = calculate_sigmas_scheduler(real_model, scheduler, steps).to(device)
+        else:
+            new_steps = int(steps / part.denoise)
+            sigmas = calculate_sigmas_scheduler(real_model, scheduler, new_steps).to(device)
+            new_sigmas[part.batch_index] = sigmas[-(steps + 1):]
+    return new_sigmas
+
+
+def generate_noise(model_wrap, sigmas, noise):
+    if max_denoise(model_wrap, sigmas):
+        n = noise * torch.sqrt(1.0 + sigmas[0] ** 2.0)
+    else:
+        n = noise * sigmas[0]
+    return n
+
+
+def get_ancestral_step(sigma_from: torch.Tensor, sigma_to: torch.Tensor, eta: float = 1.) -> (
+        torch.Tensor, torch.Tensor):
+    if not eta:
+        return sigma_to, torch.zeros_like(sigma_to)
+    sigma_up = torch.min(sigma_to,
+                         eta * (sigma_to ** 2 * (sigma_from ** 2 - sigma_to ** 2) / sigma_from ** 2) ** 0.5)
+    sigma_down = (sigma_to ** 2 - sigma_up ** 2) ** 0.5
+
+    return sigma_down, sigma_up