only remove module if necessary

2022-06-22 13:42:09 +00:00 · 2022-06-22 13:42:09 +00:00 · 40e28e8bf4
parent fc596c8625
commit 40e28e8bf4
1 changed files with 1 additions and 1 deletions
--- a/src/diffusers/pipeline_utils.py
+++ b/src/diffusers/pipeline_utils.py
@ -86,7 +86,7 @@ class DiffusionPipeline(ConfigMixin):
        model_index_dict = dict(self.config)
        model_index_dict.pop("_class_name")
        model_index_dict.pop("_diffusers_version")
-        model_index_dict.pop("_module")
+        model_index_dict.pop("_module", None)

        for pipeline_component_name in model_index_dict.keys():
            sub_model = getattr(self, pipeline_component_name)