fix: update cast_bias_weight usage and fix motion model dtype loading

- Fix #572: Use offloadable=True with cast_bias_weight in groupnorm_mm_forward
  to match ComfyUI's updated 3-return-value API, and call uncast_bias_weight
  after use for proper async-offload support.
- Fix #569: Move .to(dtype) call after load_state_dict across all model loading
  functions. Previously, the dtype cast was done before load_state_dict, which
  overwrote weights with the original float32 values from disk, causing Half vs
  Float mismatches at runtime when the UNet runs in float16. Affected functions:
  load_motion_module_gen1, load_motion_module_gen2, create_fresh_motion_module,
  create_fresh_encoder_only_model, inject_img_encoder_into_model,
  inject_pia_conv_in_into_model, inject_camera_encoder_into_model,
  load_hmreferenceadapter.

Co-authored-by: Amp <amp@ampcode.com>
Amp-Thread-ID: https://ampcode.com/threads/T-019d3bef-9580-7798-aec3-7a009a0f62a4

animatediff/adapter_hellomeme.py

@@ -70,9 +70,9 @@ def load_hmreferenceadapter(model_name: str):
     else:
         ops = comfy.ops.manual_cast
     hmref = HMReferenceAdapter(ops=ops)
-    hmref.to(comfy.model_management.unet_dtype())
     hmref.to(comfy.model_management.unet_offload_device())
     load_result = hmref.load_state_dict(state_dict, strict=True)
+    hmref.to(comfy.model_management.unet_dtype())
     hmref_model = create_HMModelPatcher(model=hmref, load_device=comfy.model_management.get_torch_device(),
                                         offload_device=comfy.model_management.unet_offload_device())
     return hmref_model

animatediff/model_injection.py

@@ -841,9 +841,9 @@ def load_motion_module_gen1(model_name: str, model: ModelPatcher, motion_lora: M
     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)
+    ad_wrapper.to(model.model_dtype())
     verify_load_result(load_result=load_result, mm_info=mm_info)
     # wrap motion_module into a ModelPatcher, to allow motion lora patches
     motion_model = create_MotionModelPatcher(model=ad_wrapper, load_device=model.load_device, offload_device=model.offload_device)
@@ -865,9 +865,9 @@ def load_motion_module_gen2(model_name: str, motion_model_settings: AnimateDiffS
     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)
+    ad_wrapper.to(comfy.model_management.unet_dtype())
     verify_load_result(load_result=load_result, mm_info=mm_info)
     # wrap motion_module into a ModelPatcher, to allow motion lora patches
     motion_model = create_MotionModelPatcher(model=ad_wrapper, load_device=comfy.model_management.get_torch_device(),
@@ -907,34 +907,34 @@ def verify_load_result(load_result: IncompatibleKeys, mm_info: AnimateDiffInfo):
 
 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())
+    ad_wrapper.to(comfy.model_management.unet_dtype())
     return create_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)
+    ad_wrapper.to(comfy.model_management.unet_dtype())
     return create_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())
+    motion_model.model.img_encoder.to(comfy.model_management.unet_dtype())
 
 
 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.conv_in.to(comfy.model_management.unet_dtype())
     motion_model.model.mm_info.mm_format = AnimateDiffFormat.PIA
 
 
@@ -956,9 +956,9 @@ def inject_camera_encoder_into_model(motion_model: MotionModelPatcher, camera_ct
     # 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.to(dtype=comfy.model_management.unet_dtype())
     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

animatediff/sampling.py

@@ -203,10 +203,13 @@ def groupnorm_mm_factory(params: InjectionParams, manual_cast=False):
 
         input = rearrange(input, "(b f) c h w -> b c f h w", b=batched_conds)
         if manual_cast:
-            weight, bias = comfy.ops.cast_bias_weight(self, input)
+            weight, bias, offload_stream = comfy.ops.cast_bias_weight(self, input, offloadable=True)
         else:
             weight, bias = self.weight, self.bias
+            offload_stream = None
         input = group_norm(input, self.num_groups, weight, bias, self.eps)
+        if offload_stream is not None:
+            comfy.ops.uncast_bias_weight(self, weight, bias, offload_stream)
         input = rearrange(input, "b c f h w -> (b f) c h w", b=batched_conds)
         return input
     return groupnorm_mm_forward
@@ -754,23 +757,6 @@ def get_resized_cond(cond_in, x_in: Tensor, full_idxs: list[int], context_length
 def sliding_calc_cond_batch(executor: Callable, model, conds: list[list[dict]], x_in: Tensor, timestep, model_options):
     ADGS: AnimateDiffGlobalState = model_options["transformer_options"]["ADGS"]
     if not ADGS.is_using_sliding_context():
-        if model_options["transformer_options"].get("context_window", None) is not None:
-            window = model_options["transformer_options"]["context_window"]
-            model_options["transformer_options"]["ad_params"]["sub_idxs"] = window.index_list
-            model_options["transformer_options"]["ad_params"]["context_length"] = len(window.index_list)
-            device = None
-            if device is None:
-                motion_models_devices = ADGS.motion_models_devices.values()
-            else:
-                motion_models_devices = ADGS.motion_models_devices.get(device, None)
-                if motion_models_devices is None:
-                    motion_models_devices = []
-                else:
-                    motion_models_devices = [motion_models_devices]
-                model = ADGS.model_patcher_devices[device].model
-            for motion_models in motion_models_devices:
-                motion_models.set_sub_idxs(window.index_list)
-                motion_models.set_video_length(len(window.index_list), ADGS.params.full_length)
         return executor(model, conds, x_in, timestep, model_options)
 
     # get context windows

pyproject.toml

@@ -1,7 +1,7 @@
 [project]
 name = "comfyui-animatediff-evolved"
 description = "Improved AnimateDiff integration for ComfyUI."
-version = "1.5.6"
+version = "1.5.7"
 license = { file = "LICENSE" }
 dependencies = []
 
@@ -13,3 +13,4 @@ Repository = "https://github.com/Kosinkadink/ComfyUI-AnimateDiff-Evolved"
 PublisherId = "kosinkadink"
 DisplayName = "ComfyUI-AnimateDiff-Evolved"
 Icon = ""
+requires-comfyui = ">=0.3.68"