Fix model offload bug when key isn't present (#5030)

* fix model offload bug when key isn't present * make style

Fix model offload bug when key isn't present (#5030)
* fix model offload bug when key isn't present * make style
e2033d2d · Dhruv Nair · GitHub · 19edca82 · e2033d2d
Unverified Commit e2033d2d authored Sep 14, 2023 by Dhruv Nair Committed by GitHub Sep 14, 2023
Hide whitespace changes
Inline Side-by-side

Showing with 1 addition and 1 deletion

src/diffusers/pipelines/pipeline_utils.py src/diffusers/pipelines/pipeline_utils.py +1 -1

No files found.
--- a/src/diffusers/pipelines/pipeline_utils.py
+++ b/src/diffusers/pipelines/pipeline_utils.py
@@ -1252,7 +1252,7 @@ class DiffusionPipeline(ConfigMixin, PushToHubMixin):
        self._all_hooks = []
        hook = None
        for model_str in self.model_cpu_offload_seq.split("->"):
-            model = all_model_components.pop(model_str)
+            model = all_model_components.pop(model_str, None)
            if not isinstance(model, torch.nn.Module):
                continue