From 81df9c85de59aae0e53ab79dbf01d024f6b9f3ba Mon Sep 17 00:00:00 2001 From: Patrick von Platen Date: Tue, 14 Nov 2023 12:08:03 +0100 Subject: [PATCH] Unwrap models everywhere (#5789) more debug --- src/diffusers/pipelines/pipeline_utils.py | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) diff --git a/src/diffusers/pipelines/pipeline_utils.py b/src/diffusers/pipelines/pipeline_utils.py index 2e5a5a20b6..674ff4d9e3 100644 --- a/src/diffusers/pipelines/pipeline_utils.py +++ b/src/diffusers/pipelines/pipeline_utils.py @@ -560,10 +560,7 @@ class DiffusionPipeline(ConfigMixin, PushToHubMixin): register_dict = {name: (None, None)} else: # register the config from the original module, not the dynamo compiled one - if is_compiled_module(module): - not_compiled_module = module._orig_mod - else: - not_compiled_module = module + not_compiled_module = _unwrap_model(module) library = not_compiled_module.__module__.split(".")[0] @@ -666,7 +663,7 @@ class DiffusionPipeline(ConfigMixin, PushToHubMixin): # Dynamo wraps the original model in a private class. # I didn't find a public API to get the original class. if is_compiled_module(sub_model): - sub_model = sub_model._orig_mod + sub_model = _unwrap_model(sub_model) model_cls = sub_model.__class__ save_method_name = None