From 5ce8e040aa45a37b40ece334b41c8ddefe6167d8 Mon Sep 17 00:00:00 2001 From: Dhruv Nair Date: Mon, 1 Jul 2024 12:32:56 +0000 Subject: [PATCH 1/2] update --- src/diffusers/models/unets/unet_motion_model.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/diffusers/models/unets/unet_motion_model.py b/src/diffusers/models/unets/unet_motion_model.py index 584014c78c6f..6c3be7652fb7 100644 --- a/src/diffusers/models/unets/unet_motion_model.py +++ b/src/diffusers/models/unets/unet_motion_model.py @@ -597,7 +597,8 @@ def from_unet2d( if not config.get("num_attention_heads"): config["num_attention_heads"] = config["attention_head_dim"] - config = FrozenDict(config) + expected_kwargs, optional_kwargs = cls._get_signature_keys(cls) + config = FrozenDict({k: config.get(k) for k in config if k in expected_kwargs or k in optional_kwargs}) model = cls.from_config(config) if not load_weights: From a12d8d90e230cebef6a68f38c730f571b4bed878 Mon Sep 17 00:00:00 2001 From: Dhruv Nair Date: Tue, 2 Jul 2024 10:20:42 +0530 Subject: [PATCH 2/2] Update src/diffusers/models/unets/unet_motion_model.py Co-authored-by: YiYi Xu --- src/diffusers/models/unets/unet_motion_model.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/diffusers/models/unets/unet_motion_model.py b/src/diffusers/models/unets/unet_motion_model.py index 6c3be7652fb7..e2657e56901f 100644 --- a/src/diffusers/models/unets/unet_motion_model.py +++ b/src/diffusers/models/unets/unet_motion_model.py @@ -599,6 +599,7 @@ def from_unet2d( expected_kwargs, optional_kwargs = cls._get_signature_keys(cls) config = FrozenDict({k: config.get(k) for k in config if k in expected_kwargs or k in optional_kwargs}) + config["_class_name"] = cls.__name__ model = cls.from_config(config) if not load_weights: