1069 lines
54 KiB
Python
1069 lines
54 KiB
Python
import copy
|
|
from typing import Union, Callable
|
|
from collections import namedtuple
|
|
|
|
from einops import rearrange
|
|
from torch import Tensor
|
|
import torch.nn.functional as F
|
|
import torch
|
|
import uuid
|
|
import math
|
|
|
|
import comfy.conds
|
|
import comfy.lora
|
|
import comfy.model_management
|
|
import comfy.utils
|
|
from comfy.model_patcher import ModelPatcher
|
|
from comfy.patcher_extension import WrappersMP, PatcherInjection
|
|
from comfy.model_base import BaseModel
|
|
from comfy.sd import CLIP, VAE
|
|
|
|
from .ad_settings import AnimateDiffSettings, AdjustPE, AdjustWeight
|
|
from .adapter_cameractrl import CameraPoseEncoder, CameraEntry, prepare_pose_embedding
|
|
from .context import ContextOptions, ContextOptions, ContextOptionsGroup
|
|
from .motion_module_ad import (AnimateDiffModel, AnimateDiffFormat, AnimateDiffInfo, EncoderOnlyAnimateDiffModel, VersatileAttention, PerBlock, AllPerBlocks,
|
|
VanillaTemporalModule, has_mid_block, normalize_ad_state_dict, get_position_encoding_max_len)
|
|
from .logger import logger
|
|
from .utils_motion import (ADKeyframe, ADKeyframeGroup, MotionCompatibilityError, InputPIA,
|
|
get_combined_multival, get_combined_input, get_combined_input_effect_multival,
|
|
ade_broadcast_image_to, extend_to_batch_size, prepare_mask_batch)
|
|
from .conditioning import HookRef, LoraHook, LoraHookGroup, LoraHookMode
|
|
from .motion_lora import MotionLoraInfo, MotionLoraList
|
|
from .utils_model import get_motion_lora_path, get_motion_model_path, get_sd_model_type, vae_encode_raw_batched
|
|
from .sample_settings import SampleSettings, SeedNoiseGeneration
|
|
|
|
|
|
class ModelPatcherHelper:
|
|
SAMPLE_SETTINGS = "ADE_sample_settings"
|
|
PARAMS = "ADE_params"
|
|
ADE = "ADE"
|
|
|
|
def __init__(self, model: ModelPatcher):
|
|
self.model = model
|
|
|
|
def set_all_properties(self, outer_sampler_wrapper: Callable, calc_cond_batch_wrapper: Callable,
|
|
params: 'InjectionParams', sample_settings: SampleSettings=None, motion_models: 'MotionModelGroup'=None):
|
|
self.set_outer_sample_wrapper(outer_sampler_wrapper)
|
|
self.set_calc_cond_batch_wrapper(calc_cond_batch_wrapper)
|
|
self.set_sample_settings(sample_settings = sample_settings if sample_settings is not None else SampleSettings())
|
|
self.set_params(params)
|
|
if motion_models is not None:
|
|
self.set_motion_models(motion_models.models.copy())
|
|
self.set_forward_timestep_embed_patch()
|
|
else:
|
|
self.remove_motion_models()
|
|
self.remove_forward_timestep_embed_patch()
|
|
|
|
def get_adgs(self):
|
|
pass
|
|
|
|
def get_motion_models(self) -> list['MotionModelPatcher']:
|
|
return self.model.additional_models.get(self.ADE, [])
|
|
|
|
def set_motion_models(self, motion_models: list['MotionModelPatcher']):
|
|
self.model.set_additional_models(self.ADE, motion_models)
|
|
self.model.set_injections(self.ADE,
|
|
[PatcherInjection(inject=inject_motion_models, eject=eject_motion_models)])
|
|
|
|
def remove_motion_models(self):
|
|
self.model.remove_additional_models(self.ADE)
|
|
self.model.remove_injections(self.ADE)
|
|
|
|
def cleanup_motion_models(self):
|
|
for motion_model in self.get_motion_models():
|
|
motion_model.cleanup()
|
|
|
|
|
|
def set_forward_timestep_embed_patch(self):
|
|
self.remove_forward_timestep_embed_patch()
|
|
self.model.set_model_forward_timestep_embed_patch(create_forward_timestep_embed_patch())
|
|
|
|
def remove_forward_timestep_embed_patch(self):
|
|
if "transformer_options" in self.model.model_options:
|
|
transformer_options = self.model.model_options["transformer_options"]
|
|
if "patches" in transformer_options:
|
|
patches = transformer_options["patches"]
|
|
if "forward_timestep_embed_patch" in patches:
|
|
forward_timestep_patches: list = patches["forward_timestep_embed_patch"]
|
|
to_remove = []
|
|
for idx, patch in enumerate(forward_timestep_patches):
|
|
if patch[1] == forward_timestep_embed_patch_ade:
|
|
to_remove.append(idx)
|
|
for idx in to_remove:
|
|
forward_timestep_patches.pop(idx)
|
|
|
|
|
|
##########################
|
|
# motion models helpers
|
|
def set_video_length(self, video_length: int, full_length: int):
|
|
for motion_model in self.get_motion_models():
|
|
motion_model.model.set_video_length(video_length=video_length, full_length=full_length)
|
|
|
|
def get_name_string(self, show_version=False):
|
|
identifiers = []
|
|
for motion_model in self.get_motion_models():
|
|
id = motion_model.model.mm_info.mm_name
|
|
if show_version:
|
|
id += f":{motion_model.model.mm_info.mm_version}"
|
|
identifiers.append(id)
|
|
return ", ".join(identifiers)
|
|
##########################
|
|
|
|
|
|
def get_sample_settings(self) -> SampleSettings:
|
|
return self.model.get_attachment(self.SAMPLE_SETTINGS)
|
|
|
|
def set_sample_settings(self, sample_settings: SampleSettings):
|
|
self.model.set_attachments(self.SAMPLE_SETTINGS, sample_settings)
|
|
|
|
|
|
def get_params(self) -> 'InjectionParams':
|
|
return self.model.get_attachment(self.PARAMS)
|
|
|
|
def set_params(self, params: 'InjectionParams'):
|
|
self.model.set_attachments(self.PARAMS, params)
|
|
|
|
|
|
def set_outer_sample_wrapper(self, wrapper: Callable):
|
|
self.model.remove_wrappers_with_key(WrappersMP.OUTER_SAMPLE, self.ADE)
|
|
self.model.add_wrapper_with_key(WrappersMP.OUTER_SAMPLE, self.ADE, wrapper)
|
|
|
|
def set_calc_cond_batch_wrapper(self, wrapper: Callable):
|
|
self.model.remove_wrappers_with_key(WrappersMP.CALC_COND_BATCH, self.ADE)
|
|
self.model.add_wrapper_with_key(WrappersMP.CALC_COND_BATCH, self.ADE, wrapper)
|
|
|
|
def remove_wrappers(self):
|
|
self.model.remove_wrappers_with_key(WrappersMP.OUTER_SAMPLE, self.ADE)
|
|
self.model.remove_wrappers_with_key(WrappersMP.CALC_COND_BATCH, self.ADE)
|
|
|
|
def pre_run(self):
|
|
# TODO: could implement this as a ModelPatcher ON_PRE_RUN callback
|
|
for motion_model in self.get_motion_models():
|
|
motion_model.pre_run(self.model)
|
|
self.get_sample_settings().pre_run(self.model)
|
|
|
|
|
|
def inject_motion_models(patcher: ModelPatcher):
|
|
helper = ModelPatcherHelper(patcher)
|
|
motion_models = helper.get_motion_models()
|
|
for mm in motion_models:
|
|
mm.model.inject(patcher)
|
|
|
|
|
|
def eject_motion_models(patcher: ModelPatcher):
|
|
helper = ModelPatcherHelper(patcher)
|
|
motion_models = helper.get_motion_models()
|
|
for mm in motion_models:
|
|
mm.model.eject(patcher)
|
|
|
|
|
|
def create_forward_timestep_embed_patch():
|
|
return (VanillaTemporalModule, forward_timestep_embed_patch_ade)
|
|
|
|
def forward_timestep_embed_patch_ade(layer, x, emb, context, transformer_options, output_shape, time_context, num_video_frames, image_only_indicator, *args, **kwargs):
|
|
return layer(x, context)
|
|
|
|
|
|
class MotionModelPatcher(ModelPatcher):
|
|
# Mostly here so that type hints work in IDEs
|
|
def __init__(self, *args, **kwargs):
|
|
super().__init__(*args, **kwargs)
|
|
self.model: AnimateDiffModel = self.model
|
|
self.timestep_percent_range = (0.0, 1.0)
|
|
self.timestep_range: tuple[float, float] = None
|
|
self.keyframes: ADKeyframeGroup = ADKeyframeGroup()
|
|
|
|
self.scale_multival: Union[float, Tensor, None] = None
|
|
self.effect_multival: Union[float, Tensor, None] = None
|
|
self.per_block_list: Union[list[PerBlock], None] = None
|
|
|
|
# AnimateLCM-I2V
|
|
self.orig_ref_drift: float = None
|
|
self.orig_insertion_weights: list[float] = None
|
|
self.orig_apply_ref_when_disabled = False
|
|
self.orig_img_latents: Tensor = None
|
|
self.img_features: list[int, Tensor] = None # temporary
|
|
self.img_latents_shape: tuple = None
|
|
|
|
# CameraCtrl
|
|
self.orig_camera_entries: list[CameraEntry] = None
|
|
self.camera_features: list[Tensor] = None # temporary
|
|
self.camera_features_shape: tuple = None
|
|
self.cameractrl_multival: Union[float, Tensor] = None
|
|
|
|
# PIA
|
|
self.orig_pia_images: Tensor = None
|
|
self.pia_vae: VAE = None
|
|
self.pia_input: InputPIA = None
|
|
self.cached_pia_c_concat: comfy.conds.CONDNoiseShape = None # cached
|
|
self.prev_pia_latents_shape: tuple = None
|
|
self.prev_current_pia_input: InputPIA = None
|
|
self.pia_multival: Union[float, Tensor] = None
|
|
|
|
# FancyVideo
|
|
self.orig_fancy_images: Tensor = None
|
|
self.fancy_vae: VAE = None
|
|
self.cached_fancy_c_concat: comfy.conds.CONDNoiseShape = None # cached
|
|
self.prev_fancy_latents_shape: tuple = None
|
|
self.fancy_multival: Union[float, Tensor] = None
|
|
|
|
# temporary variables
|
|
self.current_used_steps = 0
|
|
self.current_keyframe: ADKeyframe = None
|
|
self.current_index = -1
|
|
self.previous_t = -1
|
|
self.current_scale: Union[float, Tensor] = None
|
|
self.current_effect: Union[float, Tensor] = None
|
|
self.current_cameractrl_effect: Union[float, Tensor] = None
|
|
self.current_pia_input: InputPIA = None
|
|
self.combined_scale: Union[float, Tensor] = None
|
|
self.combined_effect: Union[float, Tensor] = None
|
|
self.combined_per_block_list: Union[float, Tensor] = None
|
|
self.combined_cameractrl_effect: Union[float, Tensor] = None
|
|
self.combined_pia_mask: Union[float, Tensor] = None
|
|
self.combined_pia_effect: Union[float, Tensor] = None
|
|
self.was_within_range = False
|
|
self.prev_sub_idxs = None
|
|
self.prev_batched_number = None
|
|
|
|
def load(self, device_to=None, lowvram_model_memory=0, *args, **kwargs):
|
|
to_return = super().load(device_to=device_to, lowvram_model_memory=lowvram_model_memory, *args, **kwargs)
|
|
if lowvram_model_memory > 0:
|
|
self._patch_lowvram_extras(device_to=device_to)
|
|
self._handle_float8_pe_tensors()
|
|
return to_return
|
|
|
|
def _patch_lowvram_extras(self, device_to=None):
|
|
# figure out the tensors (likely pe's) that should be cast to device besides just the named_modules
|
|
remaining_tensors = list(self.model.state_dict().keys())
|
|
named_modules = []
|
|
for n, _ in self.model.named_modules():
|
|
named_modules.append(n)
|
|
named_modules.append(f"{n}.weight")
|
|
named_modules.append(f"{n}.bias")
|
|
for name in named_modules:
|
|
if name in remaining_tensors:
|
|
remaining_tensors.remove(name)
|
|
|
|
for key in remaining_tensors:
|
|
self.patch_weight_to_device(key, device_to)
|
|
if device_to is not None:
|
|
comfy.utils.set_attr(self.model, key, comfy.utils.get_attr(self.model, key).to(device_to))
|
|
|
|
def _handle_float8_pe_tensors(self):
|
|
remaining_tensors = list(self.model.state_dict().keys())
|
|
pe_tensors = [x for x in remaining_tensors if '.pe' in x]
|
|
is_first = True
|
|
for key in pe_tensors:
|
|
if is_first:
|
|
is_first = False
|
|
if comfy.utils.get_attr(self.model, key).dtype not in [torch.float8_e5m2, torch.float8_e4m3fn]:
|
|
break
|
|
comfy.utils.set_attr(self.model, key, comfy.utils.get_attr(self.model, key).half())
|
|
|
|
def pre_run(self, model: ModelPatcher):
|
|
self.cleanup()
|
|
self.model.set_scale(self.scale_multival, self.per_block_list)
|
|
self.model.set_effect(self.effect_multival, self.per_block_list)
|
|
self.model.set_cameractrl_effect(self.cameractrl_multival)
|
|
if self.model.img_encoder is not None:
|
|
self.model.img_encoder.set_ref_drift(self.orig_ref_drift)
|
|
self.model.img_encoder.set_insertion_weights(self.orig_insertion_weights)
|
|
|
|
def initialize_timesteps(self, model: BaseModel):
|
|
self.timestep_range = (model.model_sampling.percent_to_sigma(self.timestep_percent_range[0]),
|
|
model.model_sampling.percent_to_sigma(self.timestep_percent_range[1]))
|
|
if self.keyframes is not None:
|
|
for keyframe in self.keyframes.keyframes:
|
|
keyframe.start_t = model.model_sampling.percent_to_sigma(keyframe.start_percent)
|
|
|
|
def prepare_current_keyframe(self, x: Tensor, t: Tensor):
|
|
curr_t: float = t[0]
|
|
# if curr_t was previous_t, then do nothing (already accounted for this step)
|
|
if curr_t == self.previous_t:
|
|
return
|
|
prev_index = self.current_index
|
|
# if met guaranteed steps, look for next keyframe in case need to switch
|
|
if self.current_keyframe is None or self.current_used_steps >= self.current_keyframe.guarantee_steps:
|
|
# if has next index, loop through and see if need to switch
|
|
if self.keyframes.has_index(self.current_index+1):
|
|
for i in range(self.current_index+1, len(self.keyframes)):
|
|
eval_kf = self.keyframes[i]
|
|
# check if start_t is greater or equal to curr_t
|
|
# NOTE: t is in terms of sigmas, not percent, so bigger number = earlier step in sampling
|
|
if eval_kf.start_t >= curr_t:
|
|
self.current_index = i
|
|
self.current_keyframe = eval_kf
|
|
self.current_used_steps = 0
|
|
# keep track of scale and effect multivals, accounting for inherit_missing
|
|
if self.current_keyframe.has_scale():
|
|
self.current_scale = self.current_keyframe.scale_multival
|
|
elif not self.current_keyframe.inherit_missing:
|
|
self.current_scale = None
|
|
if self.current_keyframe.has_effect():
|
|
self.current_effect = self.current_keyframe.effect_multival
|
|
elif not self.current_keyframe.inherit_missing:
|
|
self.current_effect = None
|
|
if self.current_keyframe.has_cameractrl_effect():
|
|
self.current_cameractrl_effect = self.current_keyframe.cameractrl_multival
|
|
elif not self.current_keyframe.inherit_missing:
|
|
self.current_cameractrl_effect = None
|
|
if self.current_keyframe.has_pia_input():
|
|
self.current_pia_input = self.current_keyframe.pia_input
|
|
elif not self.current_keyframe.inherit_missing:
|
|
self.current_pia_input = None
|
|
# if guarantee_steps greater than zero, stop searching for other keyframes
|
|
if self.current_keyframe.guarantee_steps > 0:
|
|
break
|
|
# if eval_kf is outside the percent range, stop looking further
|
|
else:
|
|
break
|
|
# if index changed, apply new combined values
|
|
if prev_index != self.current_index:
|
|
# combine model's scale and effect with keyframe's scale and effect
|
|
self.combined_scale = get_combined_multival(self.scale_multival, self.current_scale)
|
|
self.combined_effect = get_combined_multival(self.effect_multival, self.current_effect)
|
|
self.combined_cameractrl_effect = get_combined_multival(self.cameractrl_multival, self.current_cameractrl_effect)
|
|
self.combined_pia_mask = get_combined_input(self.pia_input, self.current_pia_input, x)
|
|
self.combined_pia_effect = get_combined_input_effect_multival(self.pia_input, self.current_pia_input)
|
|
# apply scale and effect
|
|
self.model.set_scale(self.combined_scale, self.per_block_list)
|
|
self.model.set_effect(self.combined_effect, self.per_block_list) # TODO: set combined_per_block_list
|
|
self.model.set_cameractrl_effect(self.combined_cameractrl_effect)
|
|
# apply effect - if not within range, set effect to 0, effectively turning model off
|
|
if curr_t > self.timestep_range[0] or curr_t < self.timestep_range[1]:
|
|
self.model.set_effect(0.0)
|
|
self.was_within_range = False
|
|
else:
|
|
# if was not in range last step, apply effect to toggle AD status
|
|
if not self.was_within_range:
|
|
self.model.set_effect(self.combined_effect, self.per_block_list)
|
|
self.was_within_range = True
|
|
# update steps current keyframe is used
|
|
self.current_used_steps += 1
|
|
# update previous_t
|
|
self.previous_t = curr_t
|
|
|
|
def prepare_alcmi2v_features(self, x: Tensor, cond_or_uncond: list[int], ad_params: dict[str], latent_format):
|
|
# if no img_encoder, done
|
|
if self.model.img_encoder is None:
|
|
return
|
|
batched_number = len(cond_or_uncond)
|
|
full_length = ad_params["full_length"]
|
|
sub_idxs = ad_params["sub_idxs"]
|
|
goal_length = x.size(0) // batched_number
|
|
# calculate img_features if needed
|
|
if (self.img_latents_shape is None or sub_idxs != self.prev_sub_idxs or batched_number != self.prev_batched_number
|
|
or x.shape[2] != self.img_latents_shape[2] or x.shape[3] != self.img_latents_shape[3]):
|
|
if sub_idxs is not None and self.orig_img_latents.size(0) >= full_length:
|
|
img_latents = comfy.utils.common_upscale(self.orig_img_latents[sub_idxs], x.shape[3], x.shape[2], 'nearest-exact', 'center').to(x.dtype).to(x.device)
|
|
else:
|
|
img_latents = comfy.utils.common_upscale(self.orig_img_latents, x.shape[3], x.shape[2], 'nearest-exact', 'center').to(x.dtype).to(x.device)
|
|
img_latents = latent_format.process_in(img_latents)
|
|
# make sure img_latents matches goal_length
|
|
if goal_length != img_latents.shape[0]:
|
|
img_latents = ade_broadcast_image_to(img_latents, goal_length, batched_number)
|
|
img_features = self.model.img_encoder(img_latents, goal_length, batched_number)
|
|
self.model.set_img_features(img_features=img_features, apply_ref_when_disabled=self.orig_apply_ref_when_disabled)
|
|
# cache values for next step
|
|
self.img_latents_shape = img_latents.shape
|
|
self.prev_sub_idxs = sub_idxs
|
|
self.prev_batched_number = batched_number
|
|
|
|
def prepare_camera_features(self, x: Tensor, cond_or_uncond: list[int], ad_params: dict[str]):
|
|
# if no camera_encoder, done
|
|
if self.model.camera_encoder is None:
|
|
return
|
|
batched_number = len(cond_or_uncond)
|
|
full_length = ad_params["full_length"]
|
|
sub_idxs = ad_params["sub_idxs"]
|
|
goal_length = x.size(0) // batched_number
|
|
# calculate camera_features if needed
|
|
if self.camera_features_shape is None or sub_idxs != self.prev_sub_idxs or batched_number != self.prev_batched_number:
|
|
# make sure there are enough camera_poses to match full_length
|
|
camera_poses = self.orig_camera_entries.copy()
|
|
if len(camera_poses) < full_length:
|
|
for i in range(full_length-len(camera_poses)):
|
|
camera_poses.append(camera_poses[-1])
|
|
if sub_idxs is not None:
|
|
camera_poses = [camera_poses[idx] for idx in sub_idxs]
|
|
# make sure camera_poses matches goal_length
|
|
if len(camera_poses) > goal_length:
|
|
camera_poses = camera_poses[:goal_length]
|
|
elif len(camera_poses) < goal_length:
|
|
# pad the camera_poses with the last element to match goal_length
|
|
for i in range(goal_length-len(camera_poses)):
|
|
camera_poses.append(camera_poses[-1])
|
|
# create encoded embeddings
|
|
b, c, h, w = x.shape
|
|
plucker_embedding = prepare_pose_embedding(camera_poses, image_width=w*8, image_height=h*8).to(dtype=x.dtype, device=x.device)
|
|
camera_embedding = self.model.camera_encoder(plucker_embedding, video_length=goal_length, batched_number=batched_number)
|
|
self.model.set_camera_features(camera_features=camera_embedding)
|
|
self.camera_features_shape = len(camera_embedding)
|
|
self.prev_sub_idxs = sub_idxs
|
|
self.prev_batched_number = batched_number
|
|
|
|
def get_pia_c_concat(self, model: BaseModel, x: Tensor) -> Tensor:
|
|
# if have cached shape, check if matches - if so, return cached pia_latents
|
|
if self.prev_pia_latents_shape is not None:
|
|
if self.prev_pia_latents_shape[0] == x.shape[0] and self.prev_pia_latents_shape[2] == x.shape[2] and self.prev_pia_latents_shape[3] == x.shape[3]:
|
|
# if mask is also the same for this timestep, then return cached
|
|
if self.prev_current_pia_input == self.current_pia_input:
|
|
return self.cached_pia_c_concat
|
|
# otherwise, adjust new mask, and create new cached_pia_c_concat
|
|
b, c, h ,w = x.shape
|
|
mask = prepare_mask_batch(self.combined_pia_mask, x.shape)
|
|
mask = extend_to_batch_size(mask, b)
|
|
# make sure to update prev_current_pia_input to know when is changed
|
|
self.prev_current_pia_input = self.current_pia_input
|
|
# TODO: handle self.combined_pia_effect eventually (feature hidden for now)
|
|
# the first index in dim=1 is the mask that needs to be updated - update in place
|
|
self.cached_pia_c_concat.cond[:, :1, :, :] = mask
|
|
return self.cached_pia_c_concat
|
|
self.prev_pia_latents_shape = None
|
|
# otherwise, x shape should be the cached pia_latents_shape
|
|
# get currently used models so they can be properly reloaded after perfoming VAE Encoding
|
|
cached_loaded_models = comfy.model_management.loaded_models(only_currently_used=True)
|
|
try:
|
|
b, c, h ,w = x.shape
|
|
usable_ref = self.orig_pia_images[:b]
|
|
# in diffusers, the image is scaled from [-1, 1] instead of default [0, 1],
|
|
# but form my testing, that blows out the images here, so I skip it
|
|
# usable_images = usable_images * 2 - 1
|
|
# resize images to latent's dims
|
|
usable_ref = usable_ref.movedim(-1,1)
|
|
usable_ref = comfy.utils.common_upscale(samples=usable_ref, width=w*self.pia_vae.downscale_ratio, height=h*self.pia_vae.downscale_ratio,
|
|
upscale_method="bilinear", crop="center")
|
|
usable_ref = usable_ref.movedim(1,-1)
|
|
# VAE encode images
|
|
logger.info("VAE Encoding PIA input images...")
|
|
usable_ref = model.process_latent_in(vae_encode_raw_batched(vae=self.pia_vae, pixels=usable_ref, show_pbar=False))
|
|
logger.info("VAE Encoding PIA input images complete.")
|
|
# make pia_latents match expected length
|
|
usable_ref = extend_to_batch_size(usable_ref, b)
|
|
self.prev_pia_latents_shape = x.shape
|
|
# now, take care of the mask
|
|
mask = prepare_mask_batch(self.combined_pia_mask, x.shape)
|
|
mask = extend_to_batch_size(mask, b)
|
|
#mask = mask.unsqueeze(1)
|
|
self.prev_current_pia_input = self.current_pia_input
|
|
if type(self.combined_pia_effect) == Tensor or not math.isclose(self.combined_pia_effect, 1.0):
|
|
real_pia_effect = self.combined_pia_effect
|
|
if type(self.combined_pia_effect) == Tensor:
|
|
real_pia_effect = extend_to_batch_size(prepare_mask_batch(self.combined_pia_effect, x.shape), b)
|
|
zero_mask = torch.zeros_like(mask)
|
|
mask = mask * real_pia_effect + zero_mask * (1.0 - real_pia_effect)
|
|
del zero_mask
|
|
zero_usable_ref = torch.zeros_like(usable_ref)
|
|
usable_ref = usable_ref * real_pia_effect + zero_usable_ref * (1.0 - real_pia_effect)
|
|
del zero_usable_ref
|
|
# cache pia c_concat
|
|
self.cached_pia_c_concat = comfy.conds.CONDNoiseShape(torch.cat([mask, usable_ref], dim=1))
|
|
return self.cached_pia_c_concat
|
|
finally:
|
|
comfy.model_management.load_models_gpu(cached_loaded_models)
|
|
|
|
def get_fancy_c_concat(self, model: BaseModel, x: Tensor) -> Tensor:
|
|
# if have cached shape, check if matches - if so, return cached fancy_latents
|
|
if self.prev_fancy_latents_shape is not None:
|
|
if self.prev_fancy_latents_shape[0] == x.shape[0] and self.prev_fancy_latents_shape[-2] == x.shape[-2] and self.prev_fancy_latents_shape[-1] == x.shape[-1]:
|
|
# TODO: if mask is also the same for this timestep, then retucn cached
|
|
return self.cached_fancy_c_concat
|
|
self.prev_fancy_latents_shape = None
|
|
# otherwise, x shape should be the cached fancy_latents_shape
|
|
# get currently used models so they can be properly reloaded after performing VAE Encoding
|
|
cached_loaded_models = comfy.model_management.loaded_models(only_currently_used=True)
|
|
try:
|
|
b, c, h, w = x.shape
|
|
usable_ref = self.orig_fancy_images[:b]
|
|
# resize images to latent's dims
|
|
usable_ref = usable_ref.movedim(-1,1)
|
|
usable_ref = comfy.utils.common_upscale(samples=usable_ref, width=w*self.fancy_vae.downscale_ratio, height=h*self.fancy_vae.downscale_ratio,
|
|
upscale_method="bilinear", crop="center")
|
|
usable_ref = usable_ref.movedim(1,-1)
|
|
# VAE encode images
|
|
logger.info("VAE Encoding FancyVideo input images...")
|
|
usable_ref: Tensor = model.process_latent_in(vae_encode_raw_batched(vae=self.fancy_vae, pixels=usable_ref, show_pbar=False))
|
|
logger.info("VAE Encoding FancyVideo input images complete.")
|
|
self.prev_fancy_latents_shape = x.shape
|
|
# TODO: experiment with indexes that aren't the first
|
|
# pad usable_ref with zeros
|
|
ref_length = usable_ref.shape[0]
|
|
pad_length = b - ref_length
|
|
zero_ref = torch.zeros([pad_length, c, h, w], dtype=usable_ref.dtype, device=usable_ref.device)
|
|
usable_ref = torch.cat([usable_ref, zero_ref], dim=0)
|
|
del zero_ref
|
|
# create mask
|
|
mask_ones = torch.ones([ref_length, 1, h, w], dtype=usable_ref.dtype, device=usable_ref.device)
|
|
mask_zeros = torch.zeros([pad_length, 1, h, w], dtype=usable_ref.dtype, device=usable_ref.device)
|
|
mask = torch.cat([mask_ones, mask_zeros], dim=0)
|
|
# TODO: experiment with mask strength
|
|
# cache fancy c_concat - ref first, then mask
|
|
self.cached_fancy_c_concat = comfy.conds.CONDNoiseShape(torch.cat([usable_ref, mask], dim=1))
|
|
return self.cached_fancy_c_concat
|
|
finally:
|
|
comfy.model_management.load_models_gpu(cached_loaded_models)
|
|
|
|
def is_pia(self):
|
|
return self.model.mm_info.mm_format == AnimateDiffFormat.PIA and self.orig_pia_images is not None
|
|
|
|
def is_fancyvideo(self):
|
|
return self.model.mm_info.mm_format == AnimateDiffFormat.FANCYVIDEO
|
|
|
|
def cleanup(self):
|
|
super().cleanup()
|
|
if self.model is not None:
|
|
self.model.cleanup()
|
|
# AnimateLCM-I2V
|
|
del self.img_features
|
|
self.img_features = None
|
|
self.img_latents_shape = None
|
|
# CameraCtrl
|
|
del self.camera_features
|
|
self.camera_features = None
|
|
self.camera_features_shape = None
|
|
# PIA
|
|
self.combined_pia_mask = None
|
|
self.combined_pia_effect = None
|
|
# Default
|
|
self.current_used_steps = 0
|
|
self.current_keyframe = None
|
|
self.current_index = -1
|
|
self.previous_t = -1
|
|
self.current_scale = None
|
|
self.current_effect = None
|
|
self.combined_scale = None
|
|
self.combined_effect = None
|
|
self.combined_per_block_list = None
|
|
self.was_within_range = False
|
|
self.prev_sub_idxs = None
|
|
self.prev_batched_number = None
|
|
|
|
def clone(self):
|
|
# normal ModelPatcher clone actions
|
|
n = MotionModelPatcher(self.model, self.load_device, self.offload_device, self.size, weight_inplace_update=self.weight_inplace_update)
|
|
n.patches = {}
|
|
for k in self.patches:
|
|
n.patches[k] = self.patches[k][:]
|
|
if hasattr(n, "patches_uuid"):
|
|
self.patches_uuid = n.patches_uuid
|
|
|
|
n.object_patches = self.object_patches.copy()
|
|
n.model_options = copy.deepcopy(self.model_options)
|
|
if hasattr(n, "model_keys"):
|
|
n.model_keys = self.model_keys
|
|
if hasattr(n, "backup"):
|
|
self.backup = n.backup
|
|
if hasattr(n, "object_patches_backup"):
|
|
self.object_patches_backup = n.object_patches_backup
|
|
# extra cloned params
|
|
n.timestep_percent_range = self.timestep_percent_range
|
|
n.timestep_range = self.timestep_range
|
|
n.keyframes = self.keyframes.clone()
|
|
n.scale_multival = self.scale_multival
|
|
n.effect_multival = self.effect_multival
|
|
# AnimateLCM-I2V
|
|
n.orig_img_latents = self.orig_img_latents
|
|
n.orig_ref_drift = self.orig_ref_drift
|
|
n.orig_insertion_weights = self.orig_insertion_weights.copy() if self.orig_insertion_weights is not None else self.orig_insertion_weights
|
|
n.orig_apply_ref_when_disabled = self.orig_apply_ref_when_disabled
|
|
# CameraCtrl
|
|
n.orig_camera_entries = self.orig_camera_entries
|
|
n.cameractrl_multival = self.cameractrl_multival
|
|
# PIA
|
|
n.orig_pia_images = self.orig_pia_images
|
|
n.pia_vae = self.pia_vae
|
|
n.pia_input = self.pia_input
|
|
n.pia_multival = self.pia_multival
|
|
return n
|
|
|
|
|
|
class MotionModelGroup:
|
|
def __init__(self, init_motion_model: MotionModelPatcher=None):
|
|
self.models: list[MotionModelPatcher] = []
|
|
if init_motion_model is not None:
|
|
if isinstance(init_motion_model, list):
|
|
for m in init_motion_model:
|
|
self.add(m)
|
|
else:
|
|
self.add(init_motion_model)
|
|
|
|
def add(self, mm: MotionModelPatcher):
|
|
# add to end of list
|
|
self.models.append(mm)
|
|
|
|
def add_to_start(self, mm: MotionModelPatcher):
|
|
self.models.insert(0, mm)
|
|
|
|
def __getitem__(self, index) -> MotionModelPatcher:
|
|
return self.models[index]
|
|
|
|
def is_empty(self) -> bool:
|
|
return len(self.models) == 0
|
|
|
|
def clone(self) -> 'MotionModelGroup':
|
|
cloned = MotionModelGroup()
|
|
for mm in self.models:
|
|
cloned.add(mm)
|
|
return cloned
|
|
|
|
def set_sub_idxs(self, sub_idxs: list[int]):
|
|
for motion_model in self.models:
|
|
motion_model.model.set_sub_idxs(sub_idxs=sub_idxs)
|
|
|
|
def set_view_options(self, view_options: ContextOptions):
|
|
for motion_model in self.models:
|
|
motion_model.model.set_view_options(view_options)
|
|
|
|
def set_video_length(self, video_length: int, full_length: int):
|
|
for motion_model in self.models:
|
|
motion_model.model.set_video_length(video_length=video_length, full_length=full_length)
|
|
|
|
def initialize_timesteps(self, model: BaseModel):
|
|
for motion_model in self.models:
|
|
motion_model.initialize_timesteps(model)
|
|
|
|
def pre_run(self, model: ModelPatcher):
|
|
for motion_model in self.models:
|
|
motion_model.pre_run(model)
|
|
|
|
def cleanup(self):
|
|
for motion_model in self.models:
|
|
motion_model.cleanup()
|
|
|
|
def prepare_current_keyframe(self, x: Tensor, t: Tensor):
|
|
for motion_model in self.models:
|
|
motion_model.prepare_current_keyframe(x=x, t=t)
|
|
|
|
def get_special_models(self):
|
|
pia_motion_models: list[MotionModelPatcher] = []
|
|
for motion_model in self.models:
|
|
if motion_model.is_pia() or motion_model.is_fancyvideo():
|
|
pia_motion_models.append(motion_model)
|
|
return pia_motion_models
|
|
|
|
def get_name_string(self, show_version=False):
|
|
identifiers = []
|
|
for motion_model in self.models:
|
|
id = motion_model.model.mm_info.mm_name
|
|
if show_version:
|
|
id += f":{motion_model.model.mm_info.mm_version}"
|
|
identifiers.append(id)
|
|
return ", ".join(identifiers)
|
|
|
|
|
|
def get_vanilla_model_patcher(m: ModelPatcher) -> ModelPatcher:
|
|
model = ModelPatcher(m.model, m.load_device, m.offload_device, m.size, weight_inplace_update=m.weight_inplace_update)
|
|
model.patches = {}
|
|
for k in m.patches:
|
|
model.patches[k] = m.patches[k][:]
|
|
|
|
model.object_patches = m.object_patches.copy()
|
|
model.model_options = copy.deepcopy(m.model_options)
|
|
if hasattr(model, "model_keys"):
|
|
model.model_keys = m.model_keys
|
|
return model
|
|
|
|
|
|
# adapted from https://github.com/guoyww/AnimateDiff/blob/main/animatediff/utils/convert_lora_safetensor_to_diffusers.py
|
|
# Example LoRA keys:
|
|
# down_blocks.0.motion_modules.0.temporal_transformer.transformer_blocks.0.attention_blocks.0.processor.to_q_lora.down.weight
|
|
# down_blocks.0.motion_modules.0.temporal_transformer.transformer_blocks.0.attention_blocks.0.processor.to_q_lora.up.weight
|
|
#
|
|
# Example model keys:
|
|
# down_blocks.0.motion_modules.0.temporal_transformer.transformer_blocks.0.attention_blocks.0.to_q.weight
|
|
#
|
|
def load_motion_lora_as_patches(motion_model: MotionModelPatcher, lora: MotionLoraInfo) -> None:
|
|
def get_version(has_midblock: bool):
|
|
return "v2" if has_midblock else "v1"
|
|
|
|
lora_path = get_motion_lora_path(lora.name)
|
|
logger.info(f"Loading motion LoRA {lora.name}")
|
|
state_dict = comfy.utils.load_torch_file(lora_path)
|
|
|
|
# remove all non-temporal keys (in case model has extra stuff in it)
|
|
for key in list(state_dict.keys()):
|
|
if "temporal" not in key:
|
|
del state_dict[key]
|
|
if len(state_dict) == 0:
|
|
raise ValueError(f"'{lora.name}' contains no temporal keys; it is not a valid motion LoRA!")
|
|
|
|
model_has_midblock = motion_model.model.mid_block != None
|
|
lora_has_midblock = has_mid_block(state_dict)
|
|
logger.info(f"Applying a {get_version(lora_has_midblock)} LoRA ({lora.name}) to a { motion_model.model.mm_info.mm_version} motion model.")
|
|
|
|
patches = {}
|
|
# convert lora state dict to one that matches motion_module keys and tensors
|
|
for key in state_dict:
|
|
# if motion_module doesn't have a midblock, skip mid_block entries
|
|
if not model_has_midblock:
|
|
if "mid_block" in key: continue
|
|
# only process lora down key (we will process up at the same time as down)
|
|
if "up." in key: continue
|
|
|
|
# get up key version of down key
|
|
up_key = key.replace(".down.", ".up.")
|
|
|
|
# adapt key to match motion_module key format - remove 'processor.', '_lora', 'down.', and 'up.'
|
|
model_key = key.replace("processor.", "").replace("_lora", "").replace("down.", "").replace("up.", "")
|
|
|
|
# motion_module keys have a '0.' after all 'to_out.' weight keys
|
|
if "to_out.0." not in model_key:
|
|
model_key = model_key.replace("to_out.", "to_out.0.")
|
|
|
|
weight_down = state_dict[key]
|
|
weight_up = state_dict[up_key]
|
|
# actual weights obtained by matrix multiplication of up and down weights
|
|
# save as a tuple, so that (Motion)ModelPatcher's calculate_weight function detects len==1, applying it correctly
|
|
patches[model_key] = (torch.mm(
|
|
comfy.model_management.cast_to_device(weight_up, weight_up.device, torch.float32),
|
|
comfy.model_management.cast_to_device(weight_down, weight_down.device, torch.float32)
|
|
),)
|
|
del state_dict
|
|
# add patches to motion ModelPatcher
|
|
motion_model.add_patches(patches=patches, strength_patch=lora.strength)
|
|
|
|
|
|
def load_motion_module_gen1(model_name: str, model: ModelPatcher, motion_lora: MotionLoraList = None, motion_model_settings: AnimateDiffSettings = None) -> MotionModelPatcher:
|
|
model_path = get_motion_model_path(model_name)
|
|
logger.info(f"Loading motion module {model_name}")
|
|
mm_state_dict = comfy.utils.load_torch_file(model_path, safe_load=True)
|
|
# TODO: check for empty state dict?
|
|
# get normalized state_dict and motion model info
|
|
mm_state_dict, mm_info = normalize_ad_state_dict(mm_state_dict=mm_state_dict, mm_name=model_name)
|
|
# check that motion model is compatible with sd model
|
|
model_sd_type = get_sd_model_type(model)
|
|
if model_sd_type != mm_info.sd_type:
|
|
raise MotionCompatibilityError(f"Motion module '{mm_info.mm_name}' is intended for {mm_info.sd_type} models, " \
|
|
+ f"but the provided model is type {model_sd_type}.")
|
|
# apply motion model settings
|
|
mm_state_dict = apply_mm_settings(model_dict=mm_state_dict, mm_settings=motion_model_settings)
|
|
# initialize AnimateDiffModelWrapper
|
|
ad_wrapper = AnimateDiffModel(mm_state_dict=mm_state_dict, mm_info=mm_info)
|
|
ad_wrapper.to(model.model_dtype())
|
|
ad_wrapper.to(model.offload_device)
|
|
load_result = ad_wrapper.load_state_dict(mm_state_dict, strict=False)
|
|
verify_load_result(load_result=load_result, mm_info=mm_info)
|
|
# wrap motion_module into a ModelPatcher, to allow motion lora patches
|
|
motion_model = MotionModelPatcher(model=ad_wrapper, load_device=model.load_device, offload_device=model.offload_device)
|
|
# load motion_lora, if present
|
|
if motion_lora is not None:
|
|
for lora in motion_lora.loras:
|
|
load_motion_lora_as_patches(motion_model, lora)
|
|
return motion_model
|
|
|
|
|
|
def load_motion_module_gen2(model_name: str, motion_model_settings: AnimateDiffSettings = None) -> MotionModelPatcher:
|
|
model_path = get_motion_model_path(model_name)
|
|
logger.info(f"Loading motion module {model_name} via Gen2")
|
|
mm_state_dict = comfy.utils.load_torch_file(model_path, safe_load=True)
|
|
# TODO: check for empty state dict?
|
|
# get normalized state_dict and motion model info (converts alternate AD models like HotshotXL into AD keys)
|
|
mm_state_dict, mm_info = normalize_ad_state_dict(mm_state_dict=mm_state_dict, mm_name=model_name)
|
|
# apply motion model settings
|
|
mm_state_dict = apply_mm_settings(model_dict=mm_state_dict, mm_settings=motion_model_settings)
|
|
# initialize AnimateDiffModelWrapper
|
|
ad_wrapper = AnimateDiffModel(mm_state_dict=mm_state_dict, mm_info=mm_info)
|
|
ad_wrapper.to(comfy.model_management.unet_dtype())
|
|
ad_wrapper.to(comfy.model_management.unet_offload_device())
|
|
load_result = ad_wrapper.load_state_dict(mm_state_dict, strict=False)
|
|
verify_load_result(load_result=load_result, mm_info=mm_info)
|
|
# wrap motion_module into a ModelPatcher, to allow motion lora patches
|
|
motion_model = MotionModelPatcher(model=ad_wrapper, load_device=comfy.model_management.get_torch_device(),
|
|
offload_device=comfy.model_management.unet_offload_device())
|
|
return motion_model
|
|
|
|
|
|
IncompatibleKeys = namedtuple('IncompatibleKeys', ['missing_keys', 'unexpected_keys'])
|
|
def verify_load_result(load_result: IncompatibleKeys, mm_info: AnimateDiffInfo):
|
|
error_msgs: list[str] = []
|
|
is_animatelcm = mm_info.mm_format==AnimateDiffFormat.ANIMATELCM
|
|
|
|
remove_missing_idxs = []
|
|
remove_unexpected_idxs = []
|
|
for idx, key in enumerate(load_result.missing_keys):
|
|
# NOTE: AnimateLCM has no pe keys in the model file, so any errors associated with missing pe keys can be ignored
|
|
if is_animatelcm and "pos_encoder.pe" in key:
|
|
remove_missing_idxs.append(idx)
|
|
# remove any keys to ignore in reverse order (to preserve idx correlation)
|
|
for idx in reversed(remove_unexpected_idxs):
|
|
load_result.unexpected_keys.pop(idx)
|
|
for idx in reversed(remove_missing_idxs):
|
|
load_result.missing_keys.pop(idx)
|
|
# copied over from torch.nn.Module.module class Module's load_state_dict func
|
|
if len(load_result.unexpected_keys) > 0:
|
|
error_msgs.insert(
|
|
0, 'Unexpected key(s) in state_dict: {}. '.format(
|
|
', '.join(f'"{k}"' for k in load_result.unexpected_keys)))
|
|
if len(load_result.missing_keys) > 0:
|
|
error_msgs.insert(
|
|
0, 'Missing key(s) in state_dict: {}. '.format(
|
|
', '.join(f'"{k}"' for k in load_result.missing_keys)))
|
|
if len(error_msgs) > 0:
|
|
raise RuntimeError('Error(s) in loading state_dict for {}:\n\t{}'.format(
|
|
mm_info.mm_name, "\n\t".join(error_msgs)))
|
|
|
|
|
|
def create_fresh_motion_module(motion_model: MotionModelPatcher) -> MotionModelPatcher:
|
|
ad_wrapper = AnimateDiffModel(mm_state_dict=motion_model.model.state_dict(), mm_info=motion_model.model.mm_info)
|
|
ad_wrapper.to(comfy.model_management.unet_dtype())
|
|
ad_wrapper.to(comfy.model_management.unet_offload_device())
|
|
ad_wrapper.load_state_dict(motion_model.model.state_dict())
|
|
return MotionModelPatcher(model=ad_wrapper, load_device=comfy.model_management.get_torch_device(),
|
|
offload_device=comfy.model_management.unet_offload_device())
|
|
|
|
|
|
def create_fresh_encoder_only_model(motion_model: MotionModelPatcher) -> MotionModelPatcher:
|
|
ad_wrapper = EncoderOnlyAnimateDiffModel(mm_state_dict=motion_model.model.state_dict(), mm_info=motion_model.model.mm_info)
|
|
ad_wrapper.to(comfy.model_management.unet_dtype())
|
|
ad_wrapper.to(comfy.model_management.unet_offload_device())
|
|
ad_wrapper.load_state_dict(motion_model.model.state_dict(), strict=False)
|
|
return MotionModelPatcher(model=ad_wrapper, load_device=comfy.model_management.get_torch_device(),
|
|
offload_device=comfy.model_management.unet_offload_device())
|
|
|
|
|
|
def inject_img_encoder_into_model(motion_model: MotionModelPatcher, w_encoder: MotionModelPatcher):
|
|
motion_model.model.init_img_encoder()
|
|
motion_model.model.img_encoder.to(comfy.model_management.unet_dtype())
|
|
motion_model.model.img_encoder.to(comfy.model_management.unet_offload_device())
|
|
motion_model.model.img_encoder.load_state_dict(w_encoder.model.img_encoder.state_dict())
|
|
|
|
|
|
def inject_pia_conv_in_into_model(motion_model: MotionModelPatcher, w_pia: MotionModelPatcher):
|
|
motion_model.model.init_conv_in(w_pia.model.state_dict())
|
|
motion_model.model.conv_in.to(comfy.model_management.unet_dtype())
|
|
motion_model.model.conv_in.to(comfy.model_management.unet_offload_device())
|
|
motion_model.model.conv_in.load_state_dict(w_pia.model.conv_in.state_dict())
|
|
motion_model.model.mm_info.mm_format = AnimateDiffFormat.PIA
|
|
|
|
|
|
def inject_camera_encoder_into_model(motion_model: MotionModelPatcher, camera_ctrl_name: str):
|
|
camera_ctrl_path = get_motion_model_path(camera_ctrl_name)
|
|
full_state_dict = comfy.utils.load_torch_file(camera_ctrl_path, safe_load=True)
|
|
camera_state_dict: dict[str, Tensor] = dict()
|
|
attention_state_dict: dict[str, Tensor] = dict()
|
|
for key in full_state_dict:
|
|
if key.startswith("encoder"):
|
|
camera_state_dict[key] = full_state_dict[key]
|
|
elif "qkv_merge" in key:
|
|
attention_state_dict[key] = full_state_dict[key]
|
|
# verify has necessary keys
|
|
if len(camera_state_dict) == 0:
|
|
raise Exception("Provided CameraCtrl model had no Camera Encoder-related keys; not a valid CameraCtrl model!")
|
|
if len(attention_state_dict) == 0:
|
|
raise Exception("Provided CameraCtrl model had no qkv_merge keys; not a valid CameraCtrl model!")
|
|
# initialize CameraPoseEncoder on motion model, and load keys
|
|
camera_encoder = CameraPoseEncoder(channels=motion_model.model.layer_channels, nums_rb=2, ops=motion_model.model.ops).to(
|
|
device=comfy.model_management.unet_offload_device(),
|
|
dtype=comfy.model_management.unet_dtype()
|
|
)
|
|
camera_encoder.load_state_dict(camera_state_dict)
|
|
camera_encoder.temporal_pe_max_len = get_position_encoding_max_len(camera_state_dict, mm_name=camera_ctrl_name, mm_format=AnimateDiffFormat.ANIMATEDIFF)
|
|
motion_model.model.set_camera_encoder(camera_encoder=camera_encoder)
|
|
# initialize qkv_merge on specific attention blocks, and load keys
|
|
for key in attention_state_dict:
|
|
key = key.strip()
|
|
# to avoid handling the same qkv_merge twice, only pay attention to the bias keys (bias+weight handled together)
|
|
if key.endswith("weight"):
|
|
continue
|
|
attr_path = key.split(".processor.qkv_merge")[0]
|
|
base_key = key.split(".bias")[0]
|
|
# first, initialize qkv_merge on model
|
|
attention_obj: VersatileAttention = comfy.utils.get_attr(motion_model.model, attr_path)
|
|
attention_obj.init_qkv_merge(ops=motion_model.model.ops)
|
|
# then, apply weights to qkv_merge
|
|
qkv_merge_state_dict = {}
|
|
qkv_merge_state_dict["weight"] = attention_state_dict[f"{base_key}.weight"]
|
|
qkv_merge_state_dict["bias"] = attention_state_dict[f"{base_key}.bias"]
|
|
attention_obj.qkv_merge.load_state_dict(qkv_merge_state_dict)
|
|
attention_obj.qkv_merge = attention_obj.qkv_merge.to(
|
|
device=comfy.model_management.unet_offload_device(),
|
|
dtype=comfy.model_management.unet_dtype()
|
|
)
|
|
|
|
|
|
def validate_model_compatibility_gen2(model: ModelPatcher, motion_model: MotionModelPatcher):
|
|
# check that motion model is compatible with sd model
|
|
model_sd_type = get_sd_model_type(model)
|
|
mm_info = motion_model.model.mm_info
|
|
if model_sd_type != mm_info.sd_type:
|
|
raise MotionCompatibilityError(f"Motion module '{mm_info.mm_name}' is intended for {mm_info.sd_type} models, " \
|
|
+ f"but the provided model is type {model_sd_type}.")
|
|
|
|
|
|
def validate_per_block_compatibility(motion_model: MotionModelPatcher, all_per_blocks: AllPerBlocks):
|
|
if all_per_blocks is None or all_per_blocks.sd_type is None:
|
|
return
|
|
mm_info = motion_model.model.mm_info
|
|
if all_per_blocks.sd_type != mm_info.sd_type:
|
|
raise Exception(f"Per-Block provided is meant for {all_per_blocks.sd_type}, but provided motion module is for {mm_info.sd_type}.")
|
|
|
|
|
|
def interpolate_pe_to_length(model_dict: dict[str, Tensor], key: str, new_length: int):
|
|
pe_shape = model_dict[key].shape
|
|
temp_pe = rearrange(model_dict[key], "(t b) f d -> t b f d", t=1)
|
|
temp_pe = F.interpolate(temp_pe, size=(new_length, pe_shape[-1]), mode="bilinear")
|
|
temp_pe = rearrange(temp_pe, "t b f d -> (t b) f d", t=1)
|
|
model_dict[key] = temp_pe
|
|
del temp_pe
|
|
|
|
|
|
def interpolate_pe_to_length_diffs(model_dict: dict[str, Tensor], key: str, new_length: int):
|
|
# TODO: fill out and try out
|
|
pe_shape = model_dict[key].shape
|
|
temp_pe = rearrange(model_dict[key], "(t b) f d -> t b f d", t=1)
|
|
temp_pe = F.interpolate(temp_pe, size=(new_length, pe_shape[-1]), mode="bilinear")
|
|
temp_pe = rearrange(temp_pe, "t b f d -> (t b) f d", t=1)
|
|
model_dict[key] = temp_pe
|
|
del temp_pe
|
|
|
|
|
|
def interpolate_pe_to_length_pingpong(model_dict: dict[str, Tensor], key: str, new_length: int):
|
|
if model_dict[key].shape[1] < new_length:
|
|
temp_pe = model_dict[key]
|
|
flipped_temp_pe = torch.flip(temp_pe[:, 1:-1, :], [1])
|
|
use_flipped = True
|
|
preview_pe = None
|
|
while model_dict[key].shape[1] < new_length:
|
|
preview_pe = model_dict[key]
|
|
model_dict[key] = torch.cat([model_dict[key], flipped_temp_pe if use_flipped else temp_pe], dim=1)
|
|
use_flipped = not use_flipped
|
|
del temp_pe
|
|
del flipped_temp_pe
|
|
del preview_pe
|
|
model_dict[key] = model_dict[key][:, :new_length]
|
|
|
|
|
|
def freeze_mask_of_pe(model_dict: dict[str, Tensor], key: str):
|
|
pe_portion = model_dict[key].shape[2] // 64
|
|
first_pe = model_dict[key][:,:1,:]
|
|
model_dict[key][:,:,pe_portion:] = first_pe[:,:,pe_portion:]
|
|
del first_pe
|
|
|
|
|
|
def freeze_mask_of_attn(model_dict: dict[str, Tensor], key: str):
|
|
attn_portion = model_dict[key].shape[0] // 2
|
|
model_dict[key][:attn_portion,:attn_portion] *= 1.5
|
|
|
|
|
|
def apply_mm_settings(model_dict: dict[str, Tensor], mm_settings: AnimateDiffSettings) -> dict[str, Tensor]:
|
|
if mm_settings is None:
|
|
return model_dict
|
|
if not mm_settings.has_anything_to_apply():
|
|
return model_dict
|
|
# first, handle PE Adjustments
|
|
for adjust_pe in mm_settings.adjust_pe.adjusts:
|
|
adjust_pe: AdjustPE
|
|
if adjust_pe.has_anything_to_apply():
|
|
already_printed = False
|
|
for key in model_dict:
|
|
if "attention_blocks" in key and "pos_encoder" in key:
|
|
# apply simple motion pe stretch, if needed
|
|
if adjust_pe.has_motion_pe_stretch():
|
|
original_length = model_dict[key].shape[1]
|
|
new_pe_length = original_length + adjust_pe.motion_pe_stretch
|
|
interpolate_pe_to_length(model_dict, key, new_length=new_pe_length)
|
|
if adjust_pe.print_adjustment and not already_printed:
|
|
logger.info(f"[Adjust PE]: PE Stretch from {original_length} to {new_pe_length}.")
|
|
# apply pe_idx_offset, if needed
|
|
if adjust_pe.has_initial_pe_idx_offset():
|
|
original_length = model_dict[key].shape[1]
|
|
model_dict[key] = model_dict[key][:, adjust_pe.initial_pe_idx_offset:]
|
|
if adjust_pe.print_adjustment and not already_printed:
|
|
logger.info(f"[Adjust PE]: Offsetting PEs by {adjust_pe.initial_pe_idx_offset}; PE length to shortens from {original_length} to {model_dict[key].shape[1]}.")
|
|
# apply has_cap_initial_pe_length, if needed
|
|
if adjust_pe.has_cap_initial_pe_length():
|
|
original_length = model_dict[key].shape[1]
|
|
model_dict[key] = model_dict[key][:, :adjust_pe.cap_initial_pe_length]
|
|
if adjust_pe.print_adjustment and not already_printed:
|
|
logger.info(f"[Adjust PE]: Capping PEs (initial) from {original_length} to {model_dict[key].shape[1]}.")
|
|
# apply interpolate_pe_to_length, if needed
|
|
if adjust_pe.has_interpolate_pe_to_length():
|
|
original_length = model_dict[key].shape[1]
|
|
interpolate_pe_to_length(model_dict, key, new_length=adjust_pe.interpolate_pe_to_length)
|
|
if adjust_pe.print_adjustment and not already_printed:
|
|
logger.info(f"[Adjust PE]: Interpolating PE length from {original_length} to {model_dict[key].shape[1]}.")
|
|
# apply final_pe_idx_offset, if needed
|
|
if adjust_pe.has_final_pe_idx_offset():
|
|
original_length = model_dict[key].shape[1]
|
|
model_dict[key] = model_dict[key][:, adjust_pe.final_pe_idx_offset:]
|
|
if adjust_pe.print_adjustment and not already_printed:
|
|
logger.info(f"[Adjust PE]: Capping PEs (final) from {original_length} to {model_dict[key].shape[1]}.")
|
|
already_printed = True
|
|
# finally, handle Weight Adjustments
|
|
for adjust_w in mm_settings.adjust_weight.adjusts:
|
|
adjust_w: AdjustWeight
|
|
if adjust_w.has_anything_to_apply():
|
|
adjust_w.mark_attrs_as_unprinted()
|
|
for key in model_dict:
|
|
# apply global weight adjustments, if needed
|
|
adjust_w.perform_applicable_ops(attr=AdjustWeight.ATTR_ALL, model_dict=model_dict, key=key)
|
|
if "attention_blocks" in key:
|
|
# apply pe change, if needed
|
|
if "pos_encoder" in key:
|
|
adjust_w.perform_applicable_ops(attr=AdjustWeight.ATTR_PE, model_dict=model_dict, key=key)
|
|
else:
|
|
# apply attn change, if needed
|
|
adjust_w.perform_applicable_ops(attr=AdjustWeight.ATTR_ATTN, model_dict=model_dict, key=key)
|
|
# apply specific attn changes, if needed
|
|
# apply attn_q change, if needed
|
|
if "to_q" in key:
|
|
adjust_w.perform_applicable_ops(attr=AdjustWeight.ATTR_ATTN_Q, model_dict=model_dict, key=key)
|
|
# apply attn_q change, if needed
|
|
elif "to_k" in key:
|
|
adjust_w.perform_applicable_ops(attr=AdjustWeight.ATTR_ATTN_K, model_dict=model_dict, key=key)
|
|
# apply attn_q change, if needed
|
|
elif "to_v" in key:
|
|
adjust_w.perform_applicable_ops(attr=AdjustWeight.ATTR_ATTN_V, model_dict=model_dict, key=key)
|
|
# apply to_out changes, if needed
|
|
elif "to_out" in key:
|
|
if key.strip().endswith("weight"):
|
|
adjust_w.perform_applicable_ops(attr=AdjustWeight.ATTR_ATTN_OUT_WEIGHT, model_dict=model_dict, key=key)
|
|
elif key.strip().endswith("bias"):
|
|
adjust_w.perform_applicable_ops(attr=AdjustWeight.ATTR_ATTN_OUT_BIAS, model_dict=model_dict, key=key)
|
|
else:
|
|
adjust_w.perform_applicable_ops(attr=AdjustWeight.ATTR_OTHER, model_dict=model_dict, key=key)
|
|
return model_dict
|
|
|
|
|
|
class InjectionParams:
|
|
def __init__(self, unlimited_area_hack: bool=False, apply_mm_groupnorm_hack: bool=True,
|
|
apply_v2_properly: bool=True) -> None:
|
|
self.full_length = None
|
|
self.unlimited_area_hack = unlimited_area_hack
|
|
self.apply_mm_groupnorm_hack = apply_mm_groupnorm_hack
|
|
self.apply_v2_properly = apply_v2_properly
|
|
self.context_options: ContextOptionsGroup = ContextOptionsGroup.default()
|
|
self.motion_model_settings = AnimateDiffSettings() # Gen1
|
|
self.sub_idxs = None # value should NOT be included in clone, so it will auto reset
|
|
|
|
def set_noise_extra_args(self, noise_extra_args: dict):
|
|
noise_extra_args["context_options"] = self.context_options.clone()
|
|
|
|
def set_context(self, context_options: ContextOptionsGroup):
|
|
self.context_options = context_options.clone() if context_options else ContextOptionsGroup.default()
|
|
|
|
def is_using_sliding_context(self) -> bool:
|
|
return self.context_options.context_length is not None
|
|
|
|
def set_motion_model_settings(self, motion_model_settings: AnimateDiffSettings): # Gen1
|
|
if motion_model_settings is None:
|
|
self.motion_model_settings = AnimateDiffSettings()
|
|
else:
|
|
self.motion_model_settings = motion_model_settings
|
|
|
|
def reset_context(self):
|
|
self.context_options = ContextOptionsGroup.default()
|
|
|
|
def clone(self) -> 'InjectionParams':
|
|
new_params = InjectionParams(
|
|
self.unlimited_area_hack, self.apply_mm_groupnorm_hack, apply_v2_properly=self.apply_v2_properly,
|
|
)
|
|
new_params.full_length = self.full_length
|
|
new_params.set_context(self.context_options)
|
|
new_params.set_motion_model_settings(self.motion_model_settings) # Gen1
|
|
return new_params
|
|
|
|
def on_model_patcher_clone(self):
|
|
return self.clone()
|