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

* fix model offload bug when key isn't present * make style
huggingface · Sep 14, 2023 · 47f2d2c · 47f2d2c
1 parent af85591
commit 47f2d2c
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/diffusers/pipelines/pipeline_utils.py b/src/diffusers/pipelines/pipeline_utils.py
@@ -1255,7 +1255,7 @@ def enable_model_cpu_offload(self, gpu_id: int = 0, device: Union[torch.device,
         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