Fix balanced offload with Cascade

vladmandic · Dec 13, 2024 · 7d7bcb9 · 7d7bcb9
1 parent 8ee5103
commit 7d7bcb9
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 1 deletion.
diff --git a/modules/intel/ipex/diffusers.py b/modules/intel/ipex/diffusers.py
@@ -1,7 +1,7 @@
 import os
 from functools import wraps, cache
 import torch
-import diffusers #0.29.1 # pylint: disable=import-error
+import diffusers # pylint: disable=import-error
 from diffusers.models.attention_processor import Attention
 
 # pylint: disable=protected-access, missing-function-docstring, line-too-long

diff --git a/modules/sd_models.py b/modules/sd_models.py
@@ -481,6 +481,8 @@ def apply_balanced_offload_to_module(pipe):
         keys = [k for k in keys if k not in exclude and not k.startswith('_')]
         for module_name, module_size in get_pipe_modules(pipe): # pylint: disable=protected-access
             module = getattr(pipe, module_name, None)
+            if module is None:
+                continue
             network_layer_name = getattr(module, "network_layer_name", None)
             device_map = getattr(module, "balanced_offload_device_map", None)
             max_memory = getattr(module, "balanced_offload_max_memory", None)