Switch device types

AppMana · Aug 21, 2024 · 3ec046c · 3ec046c
1 parent 2c8acfd
commit 3ec046c
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/comfyui_ella/model.py b/comfyui_ella/model.py
@@ -314,10 +314,10 @@ def forward(self, timesteps: torch.Tensor, t5_embeds: torch.Tensor, **kwargs):
 
 class ELLA:
     def __init__(self, path: str, **kwargs) -> None:
-        self.load_device = model_management.text_encoder_device()
-        self.offload_device = model_management.text_encoder_offload_device()
-        self.dtype = model_management.text_encoder_dtype(self.load_device)
-        self.output_device = model_management.intermediate_device()
+        self.load_device = model_management.get_torch_device()
+        self.offload_device = model_management.unet_offload_device()
+        self.dtype = model_management.unet_dtype(self.load_device)
+        self.output_device = model_management.get_torch_device()
         self.model = ELLAModel()
         load_model(self.model, path, strict=True)
         self.model.to(self.dtype)  # type: ignore