From e2033d2dff6eb86b06f03e0405938a43a2608044 Mon Sep 17 00:00:00 2001
From: Dhruv Nair <dhruv.nair@gmail.com>
Date: Thu, 14 Sep 2023 14:32:06 +0530
Subject: [PATCH] Fix model offload bug when key isn't present (#5030)

* fix model offload bug when key isn't present

* make style
---
 src/diffusers/pipelines/pipeline_utils.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/diffusers/pipelines/pipeline_utils.py b/src/diffusers/pipelines/pipeline_utils.py
index 8b12ab6a0d..63e25fdf75 100644
--- 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