ostris/ai-toolkit
1
0
Fork 0
mirror of https://github.com/ostris/ai-toolkit.git synced 2026-05-01 03:31:35 +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: if self.adapter_config is not None:
self.setup_adapter() self.setup_adapter()
flush() 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 not self.is_fine_tuning:
if self.network_config is not None: if self.network_config is not None:
# TODO should we completely switch to LycorisSpecialNetwork? # TODO should we completely switch to LycorisSpecialNetwork?
@@ -1334,6 +1325,15 @@ class BaseSDTrainProcess(BaseTrainProcess):
) )
self.lr_scheduler = lr_scheduler 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() flush()
### HOOK ### ### HOOK ###
self.hook_before_train_loop() self.hook_before_train_loop()

2
toolkit/config_modules.py
View File

@@ -127,7 +127,7 @@ class NetworkConfig:
self.conv = 4 self.conv = 4
AdapterTypes = Literal['t2i', 'ip', 'ip+'] AdapterTypes = Literal['t2i', 'ip', 'ip+', 'clip', 'ilora']
class AdapterConfig: class AdapterConfig:

3
toolkit/ip_adapter.py
View File

@@ -288,6 +288,9 @@ class IPAdapter(torch.nn.Module):
output_dim=sd.unet.config['cross_attention_dim'], output_dim=sd.unet.config['cross_attention_dim'],
ff_mult=4 ff_mult=4
) )
elif adapter_config.type == 'ilora':
# we apply the clip encodings to the LoRA
image_proj_model = None
else: else:
raise ValueError(f"unknown adapter type: {adapter_config.type}") raise ValueError(f"unknown adapter type: {adapter_config.type}")