ostris/ai-toolkit
1
0
Fork 0
mirror of https://github.com/ostris/ai-toolkit.git synced 2026-01-26 16:39:47 +00:00
Code Issues Packages Projects Releases Wiki Activity

Prepwork for ilora

Browse Source
This commit is contained in:
Jaret Burkett
2024-01-12 06:41:15 -07:00
parent 290393f7ae
commit e190fbaeb8
3 changed files with 13 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files

18
jobs/process/BaseSDTrainProcess.py
View File

@@ -1119,15 +1119,6 @@ class BaseSDTrainProcess(BaseTrainProcess):
if self.adapter_config is not None:
self.setup_adapter()
flush()
### HOOk ###
self.before_dataset_load()
# load datasets if passed in the root process
if self.datasets is not None:
self.data_loader = get_dataloader_from_datasets(self.datasets, self.train_config.batch_size, self.sd)
if self.datasets_reg is not None:
self.data_loader_reg = get_dataloader_from_datasets(self.datasets_reg, self.train_config.batch_size,
self.sd)
if not self.is_fine_tuning:
if self.network_config is not None:
# TODO should we completely switch to LycorisSpecialNetwork?
@@ -1334,6 +1325,15 @@ class BaseSDTrainProcess(BaseTrainProcess):
)
self.lr_scheduler = lr_scheduler
### HOOk ###
self.before_dataset_load()
# load datasets if passed in the root process
if self.datasets is not None:
self.data_loader = get_dataloader_from_datasets(self.datasets, self.train_config.batch_size, self.sd)
if self.datasets_reg is not None:
self.data_loader_reg = get_dataloader_from_datasets(self.datasets_reg, self.train_config.batch_size,
self.sd)
flush()
### HOOK ###
self.hook_before_train_loop()