Fix OOMs happening in some cases.

A cloned model patcher sometimes reported a model was loaded on a device when it wasn't.
2024-08-06 13:27:48 -04:00
parent de17a9755e
commit b334605a66
4 changed files with 17 additions and 11 deletions
--- a/comfy/model_base.py
+++ b/comfy/model_base.py
@@ -74,6 +74,7 @@ class BaseModel(torch.nn.Module):
        self.latent_format = model_config.latent_format
        self.model_config = model_config
        self.manual_cast_dtype = model_config.manual_cast_dtype
+        self.device = device

        if not unet_config.get("disable_unet_model_creation", False):
            if self.manual_cast_dtype is not None: