From 142fddfe47074929cb1c79f8bec49facff37c7b1 Mon Sep 17 00:00:00 2001
From: Vladimir Mandic <mandic00@live.com>
Date: Tue, 22 Oct 2024 20:31:35 -0400
Subject: [PATCH] cleanup

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

diff --git a/modules/model_omnigen.py b/modules/model_omnigen.py
index 9212193dd..64c99ddd7 100644
--- a/modules/model_omnigen.py
+++ b/modules/model_omnigen.py
@@ -13,6 +13,7 @@ def load_omnigen(checkpoint_info, diffusers_load_config={}): # pylint: disable=u
     # init
     pipe.device = devices.device
     pipe.dtype = devices.dtype
+    pipe.model.device = devices.device
     pipe.separate_cfg_infer = True
     pipe.use_kv_cache = False
     pipe.model.to(device=devices.device, dtype=devices.dtype)