ostris/ai-toolkit
1
0
Fork 0
mirror of https://github.com/ostris/ai-toolkit.git synced 2026-04-29 10:41:28 +00:00
Code Issues Packages Projects Releases Wiki Activity

Reworked automagic optimizer and did more testing. Starting to really like it. Working well.

Browse Source
This commit is contained in:
Jaret Burkett
2025-04-28 08:01:10 -06:00
parent 88b3fbae37
commit 2b4c525489
5 changed files with 149 additions and 61 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
toolkit/models/base_model.py
View File

@@ -113,15 +113,15 @@ class BaseModel:
):
self.accelerator = get_accelerator()
self.custom_pipeline = custom_pipeline
self.device = str(self.accelerator.device)
self.device = device
self.dtype = dtype
self.torch_dtype = get_torch_dtype(dtype)
self.device_torch = self.accelerator.device
self.device_torch = torch.device(device)
self.vae_device_torch = self.accelerator.device
self.vae_device_torch = torch.device(device)
self.vae_torch_dtype = get_torch_dtype(model_config.vae_dtype)
self.te_device_torch = self.accelerator.device
self.te_device_torch = torch.device(device)
self.te_torch_dtype = get_torch_dtype(model_config.te_dtype)
self.model_config = model_config