Ensure models are on the correct device more reliably in BlehRefinerAfter

This commit is contained in:
blepping
2024-04-03 11:12:05 -06:00
parent 03d7d2cbf1
commit 14df853670
+5 -5
View File
@@ -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,