From 76a0a632434604f012c90372802db6813de26b0b Mon Sep 17 00:00:00 2001
From: Vladimir Mandic <mandic00@live.com>
Date: Tue, 29 Oct 2024 13:53:21 -0400
Subject: [PATCH] offload exclude lsit

Signed-off-by: Vladimir Mandic <mandic00@live.com>
---
 modules/sd_models.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/modules/sd_models.py b/modules/sd_models.py
index 7aa70b90f..734f4ab57 100644
--- a/modules/sd_models.py
+++ b/modules/sd_models.py
@@ -375,6 +375,9 @@ def set_diffuser_offload(sd_model, op: str = 'model'):
 def apply_balanced_offload(sd_model):
     from accelerate import infer_auto_device_map, dispatch_model
     from accelerate.hooks import add_hook_to_module, remove_hook_from_module, ModelHook
+    excluded = ['OmniGenPipeline']
+    if sd_model.__class__.__name__ in excluded:
+        return sd_model
 
     class dispatch_from_cpu_hook(ModelHook):
         def init_hook(self, module):