github/ai-toolkit
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

Bug fixes. Added IP adapter training for Pixart

Browse Source
This commit is contained in:
Jaret Burkett
2024-02-17 10:06:57 -07:00
parent 93b52932c1
commit 2478554c95
4 changed files with 278 additions and 49 deletions
Download Patch File Download Diff File Expand all files Collapse all files

13
toolkit/stable_diffusion_model.py
View File

@@ -1567,10 +1567,15 @@ class StableDiffusion:
named_params = self.named_parameters(vae=False, unet=unet, text_encoder=False, state_dict_keys=True)
unet_lr = unet_lr if unet_lr is not None else default_lr
params = []
for key, diffusers_key in ldm_diffusers_keymap.items():
if diffusers_key in named_params and diffusers_key not in DO_NOT_TRAIN_WEIGHTS:
if named_params[diffusers_key].requires_grad:
params.append(named_params[diffusers_key])
if self.is_pixart:
for param in named_params.values():
if param.requires_grad:
params.append(param)
else:
for key, diffusers_key in ldm_diffusers_keymap.items():
if diffusers_key in named_params and diffusers_key not in DO_NOT_TRAIN_WEIGHTS:
if named_params[diffusers_key].requires_grad:
params.append(named_params[diffusers_key])
param_data = {"params": params, "lr": unet_lr}
trainable_parameters.append(param_data)
print(f"Found {len(params)} trainable parameter in unet")