From 14df85367006c171bbd821f857da53c262cb998e Mon Sep 17 00:00:00 2001
From: blepping <blepping@tuta.io>
Date: Wed, 3 Apr 2024 11:12:05 -0600
Subject: [PATCH] Ensure models are on the correct device more reliably in
 BlehRefinerAfter

---
 py/refinerAfter.py | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/py/refinerAfter.py b/py/refinerAfter.py
index dd3e1cf..07cec0c 100644
--- a/py/refinerAfter.py
+++ b/py/refinerAfter.py
@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-import comfy
+import comfy.model_management as mm
 
 
 class BlehRefinerAfter:
@@ -34,10 +34,10 @@ class BlehRefinerAfter:
 
     @staticmethod
     def load_if_needed(model: object) -> bool:
-        if model.current_device != model.load_device:
-            comfy.model_management.load_models_gpu([model])
-            return True
-        return False
+        if mm.LoadedModel(model) in mm.current_loaded_models:
+            return False
+        mm.load_models_gpu([model])
+        return True
 
     def patch(  # noqa: PLR0911
         self,