Got CameraCtrl working, fixed weird edge case causing AD device mismatches after switching models + prompts

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.

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):

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