From 9d2d4e747bc24efefdf0808ffe2f3a770350059d Mon Sep 17 00:00:00 2001 From: Jedrzej Kosinski Date: Fri, 19 Apr 2024 13:13:33 -0500 Subject: [PATCH] Got CameraCtrl working, fixed weird edge case causing AD device mismatches after switching models + prompts --- animatediff/adapter_cameractrl.py | 71 ++++++++++++++++--------------- animatediff/model_injection.py | 38 +++++++---------- animatediff/motion_module_ad.py | 28 ++++++------ 3 files changed, 68 insertions(+), 69 deletions(-) diff --git a/animatediff/adapter_cameractrl.py b/animatediff/adapter_cameractrl.py index 43384cf..991565f 100644 --- a/animatediff/adapter_cameractrl.py +++ b/animatediff/adapter_cameractrl.py @@ -1,17 +1,15 @@ - # Modified from https://github.com/hehao13/CameraCtrl/blob/main/cameractrl/models/pose_adaptor.py # (whose parts were also taken from https://github.com/TencentARC/T2I-Adapter) import torch import torch.nn as nn import numpy as np from torch import Tensor -from collections import OrderedDict from einops import rearrange import comfy.ops from .motion_module_ad import TemporalTransformerBlock - +from .logger import logger def conv_nd(dims, *args, **kwargs): """ @@ -39,25 +37,10 @@ def avg_pool_nd(dims, *args, **kwargs): raise ValueError(f"unsupported dimensions: {dims}") -# class PoseAdapter(nn.Module): -# def __init__(self, unet, pose_encoder): -# super().__init__() -# self.unet = unet -# self.pose_encoder = pose_encoder - -# def forward(self, noisy_latents: Tensor, timesteps, encoder_hidden_states: Tensor, pose_embedding): -# # original code needed to convert from 4 dims (bf c h w) to 5 dims (b c f h w), -# # but ComfyUI already deals with everything in 4 dims -# pose_embedding_features = self.pose_encoder(pose_embedding) -# noise_pred = self.unet(noisy_latents, -# timesteps, -# encoder_hidden_states, -# pose_embedding_features).sample -# return noise_pred - class CameraEntry: # TODO: comment on each value with an explanation - def __init__(self, entry): + def __init__(self, entry: list[Tensor]): + self.entry = entry.copy() fx, fy, cx, cy = entry[1:5] self.fx = fx self.fy = fy @@ -68,6 +51,9 @@ class CameraEntry: w2c_mat_4x4[:3, :] = w2c_mat self.w2c_mat = w2c_mat_4x4 self.c2w_mat = np.linalg.inv(w2c_mat_4x4) + + def clone(self): + return CameraEntry(entry=self.entry) def get_parameter_dtype(parameter: torch.nn.Module): @@ -80,6 +66,16 @@ def get_parameter_dtype(parameter: torch.nn.Module): return buffers[0].dtype +def get_parameter_device(parameter: torch.nn.Module): + params = tuple(parameter.parameters()) + if len(params) > 0: + return params[0].device + + buffers = tuple(parameter.buffers()) + if len(buffers) > 0: + return buffers[0].device + + def custom_meshgrid(*args): # ref: https://pytorch.org/docs/stable/generated/torch.meshgrid.html?highlight=meshgrid#torch.meshgrid return torch.meshgrid(*args, indexing='ij') @@ -136,6 +132,8 @@ def ray_condition(K: Tensor, c2w: Tensor, H, W, device): def prepare_pose_embedding(cam_params: list[CameraEntry], image_width, image_height, original_pose_width=1280, original_pose_height=720): + # clone each CameraEntry in list so that CameraEntries don't get spoiled after a single run + cam_params = [entry.clone() for entry in cam_params] sample_wh_ratio = image_width / image_height pose_wh_ratio = original_pose_width / original_pose_height if pose_wh_ratio > sample_wh_ratio: @@ -156,8 +154,6 @@ def prepare_pose_embedding(cam_params: list[CameraEntry], image_width, image_hei c2ws = get_relative_pose(cam_params) c2ws = torch.as_tensor(c2ws)[None] # [1, n_frame, 4, 4] plucker_embedding = ray_condition(K, c2ws, image_height, image_width, device='cpu')[0].permute(0, 3, 1, 2).contiguous() # V, 6, H, W - #plucker_embedding = plucker_embedding[None] # B V 6 H W - #plucker_embedding = rearrange(plucker_embedding, "b f c h w -> b c f h w") plucker_embedding = rearrange(plucker_embedding, "f c h w -> c f h w") return plucker_embedding @@ -192,19 +188,19 @@ class CameraPoseEncoder(nn.Module): if j == 0 and i != 0: in_dim = channels[i - 1] out_dim = int(channels[i] / compression_factor) - conv_layer = ResnetBlock(in_dim, out_dim, down=True, ksize=ksize, sk=sk, use_conv=use_conv, ops=ops) + conv_layer = ResnetBlockCameraCtrl(in_dim, out_dim, down=True, ksize=ksize, sk=sk, use_conv=use_conv, ops=ops) elif j == 0: in_dim = channels[0] out_dim = int(channels[i] / compression_factor) - conv_layer = ResnetBlock(in_dim, out_dim, down=False, ksize=ksize, sk=sk, use_conv=use_conv, ops=ops) + conv_layer = ResnetBlockCameraCtrl(in_dim, out_dim, down=False, ksize=ksize, sk=sk, use_conv=use_conv, ops=ops) elif j == nums_rb - 1: in_dim = channels[i] / compression_factor out_dim = channels[i] - conv_layer = ResnetBlock(in_dim, out_dim, down=False, ksize=ksize, sk=sk, use_conv=use_conv, ops=ops) + conv_layer = ResnetBlockCameraCtrl(in_dim, out_dim, down=False, ksize=ksize, sk=sk, use_conv=use_conv, ops=ops) else: in_dim = int(channels[i] / compression_factor) out_dim = int(channels[i] / compression_factor) - conv_layer = ResnetBlock(in_dim, out_dim, down=False, ksize=ksize, sk=sk, use_conv=use_conv, ops=ops) + conv_layer = ResnetBlockCameraCtrl(in_dim, out_dim, down=False, ksize=ksize, sk=sk, use_conv=use_conv, ops=ops) temporal_attention_layer = TemporalTransformerBlock(dim=out_dim, num_attention_heads=temporal_attention_nhead, attention_head_dim=int(out_dim / temporal_attention_nhead), @@ -213,6 +209,7 @@ class CameraPoseEncoder(nn.Module): cross_attention_dim=None, temporal_pe=temporal_position_encoding, temporal_pe_max_len=temporal_position_encoding_max_len, + rearrange_hidden_shapes=False, # different from AD ops=ops) conv_layers.append(conv_layer) temporal_attention_layers.append(temporal_attention_layer) @@ -222,11 +219,14 @@ class CameraPoseEncoder(nn.Module): def forward(self, x: Tensor, video_length: int, batched_number: int=1): # rearrange to match expected format x = rearrange(x, "c f h w -> f c h w") + logger.info(f"x: {x.shape}, {float(x[0][0][0][-1])}") # unshuffle x = self.unshuffle(x) # extract features features = [] - x = self.encoder_conv_in(x.to(dtype=get_parameter_dtype(self))) + logger.warn(f"x dtype: {x.dtype}, device: {x.device}") + logger.warn(f"dtype: {get_parameter_dtype(self)}, device: {get_parameter_device(self)}") + x = self.encoder_conv_in(x.to(dtype=get_parameter_dtype(self), device=get_parameter_device(self))) for res_block, attention_block in zip(self.encoder_down_conv_blocks, self.encoder_down_attention_blocks): for res_layer, attention_layer in zip(res_block, attention_block): x = res_layer(x) @@ -235,14 +235,17 @@ class CameraPoseEncoder(nn.Module): x = attention_layer(x, video_length=video_length) x = rearrange(x, '(h w) b c -> b c h w', h=h, w=w) features.append(x) - for idx, x in enumerate(features): - x = rearrange(x, 'b c h w -> (h w) b c') - x = torch.cat([x] * batched_number, dim=0) - features[idx] = x + for idx, feature in enumerate(features): + logger.info(f"{idx}: {feature.shape}, {float(feature[0][0][0][0])}") + for idx, x1 in enumerate(features): + x1 = x1.to(x.dtype).to(x.device) + x1 = rearrange(x1, 'b c h w -> (h w) b c') + x1 = torch.cat([x1] * batched_number, dim=0) + features[idx] = x1 return features -class ResnetBlock(nn.Module): +class ResnetBlockCameraCtrl(nn.Module): def __init__(self, in_c, out_c, down: bool, ksize=3, sk=False, use_conv=True, ops=comfy.ops.disable_weight_init): super().__init__() @@ -261,7 +264,7 @@ class ResnetBlock(nn.Module): self.down = down if self.down == True: - self.down_opt = Downsample(in_c, use_conv=use_conv) + self.down_opt = DownsampleCameraCtrl(in_c, use_conv=use_conv) def forward(self, x: Tensor): if self.down == True: @@ -278,7 +281,7 @@ class ResnetBlock(nn.Module): return h + x -class Downsample(nn.Module): +class DownsampleCameraCtrl(nn.Module): """ A downsampling layer with an optional convolution. :param channels: channels in the inputs and outputs. diff --git a/animatediff/model_injection.py b/animatediff/model_injection.py index 3854e50..2b6ae97 100644 --- a/animatediff/model_injection.py +++ b/animatediff/model_injection.py @@ -50,12 +50,6 @@ class ModelPatcherAndInjector(ModelPatcher): def model_patches_to(self, device): super().model_patches_to(device) - if self.motion_models is not None: - for motion_model in self.motion_models.models: - try: - motion_model.model.to(device) - except Exception: - pass def patch_model(self, device_to=None, patch_weights=True): # first, perform model patching @@ -80,19 +74,11 @@ class ModelPatcherAndInjector(ModelPatcher): if self.motion_models is not None: for motion_model in self.motion_models.models: motion_model.model.inject(self) - try: - motion_model.model.to(device_to) - except Exception: - pass def eject_model(self, device_to=None): if self.motion_models is not None: for motion_model in self.motion_models.models: motion_model.model.eject(self) - try: - motion_model.model.to(device_to) - except Exception: - pass def clone(self): cloned = ModelPatcherAndInjector(self) @@ -246,10 +232,13 @@ class MotionModelPatcher(ModelPatcher): 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: - if sub_idxs is not None and len(self.orig_camera_entries) >= full_length: - camera_poses = [self.orig_camera_entries[idx] for idx in sub_idxs] - else: - camera_poses = self.orig_camera_entries.copy() + # 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 and len(self.orig_camera_entries) < full_length: + camera_poses = [camera_poses[idx] for idx in sub_idxs] # make sure camera_poses matches goal_length if goal_length != len(camera_poses): if len(camera_poses) > goal_length: @@ -551,11 +540,10 @@ def inject_camera_encoder_into_model(motion_model: MotionModelPatcher, camera_ct 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 = 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) motion_model.model.set_camera_encoder(camera_encoder=camera_encoder) # initialize qkv_merge on specific attention blocks, and load keys @@ -575,6 +563,10 @@ def inject_camera_encoder_into_model(motion_model: MotionModelPatcher, camera_ct 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): diff --git a/animatediff/motion_module_ad.py b/animatediff/motion_module_ad.py index be99036..eb6c1ae 100644 --- a/animatediff/motion_module_ad.py +++ b/animatediff/motion_module_ad.py @@ -685,6 +685,7 @@ class TemporalTransformer3DModel(nn.Module): cross_frame_attention_mode=None, temporal_pe=False, temporal_pe_max_len=24, + rearrange_hidden_shapes=True, ops=comfy.ops.disable_weight_init, ): super().__init__() @@ -719,6 +720,7 @@ class TemporalTransformer3DModel(nn.Module): cross_frame_attention_mode=cross_frame_attention_mode, temporal_pe=temporal_pe, temporal_pe_max_len=temporal_pe_max_len, + rearrange_hidden_shapes=rearrange_hidden_shapes, ops=ops, ) for d in range(num_layers) @@ -847,6 +849,7 @@ class TemporalTransformerBlock(nn.Module): cross_frame_attention_mode=None, temporal_pe=False, temporal_pe_max_len=24, + rearrange_hidden_shapes=True, ops=comfy.ops.disable_weight_init, ): super().__init__() @@ -870,6 +873,7 @@ class TemporalTransformerBlock(nn.Module): cross_frame_attention_mode=cross_frame_attention_mode, temporal_pe=temporal_pe, temporal_pe_max_len=temporal_pe_max_len, + rearrange_hidden_shapes=rearrange_hidden_shapes, ops=ops, ) ) @@ -1018,6 +1022,7 @@ class VersatileAttention(CrossAttentionMM): cross_frame_attention_mode=None, temporal_pe=False, temporal_pe_max_len=24, + rearrange_hidden_shapes=True, ops=comfy.ops.disable_weight_init, *args, **kwargs, @@ -1031,6 +1036,7 @@ class VersatileAttention(CrossAttentionMM): self.query_dim: int = kwargs["query_dim"] self.qkv_merge: comfy.ops.disable_weight_init.Linear = None self.camera_feature_enabled = False + self.rearrange_hidden_shapes = rearrange_hidden_shapes self.pos_encoder = ( PositionalEncoding( @@ -1070,10 +1076,11 @@ class VersatileAttention(CrossAttentionMM): if self.attention_mode != "Temporal": raise NotImplementedError - d = hidden_states.shape[1] - hidden_states = rearrange( - hidden_states, "(b f) d c -> (b d) f c", f=video_length - ) + if self.rearrange_hidden_shapes: + d = hidden_states.shape[1] + hidden_states = rearrange( + hidden_states, "(b f) d c -> (b d) f c", f=video_length + ) if self.pos_encoder is not None: hidden_states = self.pos_encoder(hidden_states).to(hidden_states.dtype) @@ -1085,13 +1092,9 @@ class VersatileAttention(CrossAttentionMM): ) if self.camera_feature_enabled and self.qkv_merge is not None and mm_kwargs is not None and "camera_feature" in mm_kwargs: - logger.info("performing camera_feature") - camera_feature = mm_kwargs["camera_feature"] - #camera_feature = rearrange(camera_feature, "b c h w -> (h w) b c") - #camera_feature = torch.cat([camera_feature] * (camera_feature.shape[1] // video_length), dim=0) - #camera_feature = rearrange(camera_feature, "d (b f) c -> (b d) f c", f=video_length) - #camera_feature = rearrange(camera_feature, "d (b f) c -> (d b) f c", f=video_length) - hidden_states = self.qkv_merge(hidden_states + camera_feature) + hidden_states + camera_feature: Tensor = mm_kwargs["camera_feature"] + cameractrl_effect = 1.0 + hidden_states = (self.qkv_merge(hidden_states + camera_feature) + hidden_states) * cameractrl_effect + hidden_states * (1. - cameractrl_effect) hidden_states = super().forward( hidden_states, @@ -1101,7 +1104,8 @@ class VersatileAttention(CrossAttentionMM): scale_mask=scale_mask, ) - hidden_states = rearrange(hidden_states, "(b d) f c -> (b f) d c", d=d) + if self.rearrange_hidden_shapes: + hidden_states = rearrange(hidden_states, "(b d) f c -> (b f) d c", d=d) return hidden_states