Skip to content

GitLab

Unverified Commit 81df9c85 authored by Patrick von Platen's avatar Patrick von Platen Committed by GitHub
Browse files

Unwrap models everywhere (#5789) 

more debug
parent bfe94a39
Hide whitespace changes
Inline Side-by-side
Showing with 2 additions and 5 deletions
src/diffusers/pipelines/pipeline_utils.py
View file @ 81df9c85
...@@ -560,10 +560,7 @@ class DiffusionPipeline(ConfigMixin, PushToHubMixin): ...@@ -560,10 +560,7 @@ class DiffusionPipeline(ConfigMixin, PushToHubMixin):
register_dict = {name: (None, None)} register_dict = {name: (None, None)}
else: else:
# register the config from the original module, not the dynamo compiled one # register the config from the original module, not the dynamo compiled one
if is_compiled_module(module): not_compiled_module = _unwrap_model(module)
not_compiled_module = module._orig_mod
else:
not_compiled_module = module
library = not_compiled_module.__module__.split(".")[0] library = not_compiled_module.__module__.split(".")[0]
...@@ -666,7 +663,7 @@ class DiffusionPipeline(ConfigMixin, PushToHubMixin): ...@@ -666,7 +663,7 @@ class DiffusionPipeline(ConfigMixin, PushToHubMixin):
# Dynamo wraps the original model in a private class. # Dynamo wraps the original model in a private class.
# I didn't find a public API to get the original class. # I didn't find a public API to get the original class.
if is_compiled_module(sub_model): if is_compiled_module(sub_model):
sub_model = sub_model._orig_mod sub_model = _unwrap_model(sub_model)
model_cls = sub_model.__class__ model_cls = sub_model.__class__
save_method_name = None save_method_name = None
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment