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

Adapter work. Bug fixes. Auto adjust LR when resuming optimizer.

Browse Source
This commit is contained in:
Jaret Burkett
2024-03-17 10:21:47 -06:00
parent 72de68d8aa
commit 016687bda1
8 changed files with 84 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
toolkit/config_modules.py
View File

@@ -218,7 +218,7 @@ class TrainConfig:
self.xformers = kwargs.get('xformers', False)
self.sdp = kwargs.get('sdp', False)
self.train_unet = kwargs.get('train_unet', True)
self.train_text_encoder = kwargs.get('train_text_encoder', True)
self.train_text_encoder = kwargs.get('train_text_encoder', False)
self.train_refiner = kwargs.get('train_refiner', True)
self.train_turbo = kwargs.get('train_turbo', False)
self.show_turbo_outputs = kwargs.get('show_turbo_outputs', False)
@@ -298,6 +298,9 @@ class TrainConfig:
self.do_random_cfg = kwargs.get('do_random_cfg', False)
self.cfg_scale = kwargs.get('cfg_scale', 1.0)
self.max_cfg_scale = kwargs.get('max_cfg_scale', self.cfg_scale)
self.cfg_rescale = kwargs.get('cfg_rescale', None)
if self.cfg_rescale is None:
self.cfg_rescale = self.cfg_scale
# applies the inverse of the prediction mean and std to the target to correct
# for norm drift

18
toolkit/custom_adapter.py
View File

@@ -132,9 +132,16 @@ class CustomAdapter(torch.nn.Module):
vision_tokens = ((self.vision_encoder.config.image_size // self.vision_encoder.config.patch_size) ** 2)
if self.config.image_encoder_arch == 'clip':
vision_tokens = vision_tokens + 1
vision_hidden_size = self.vision_encoder.config.hidden_size
if self.config.clip_layer == 'image_embeds':
vision_tokens = 1
vision_hidden_size = self.vision_encoder.config.projection_dim
self.ilora_module = InstantLoRAModule(
vision_tokens=vision_tokens,
vision_hidden_size=self.vision_encoder.config.hidden_size,
vision_hidden_size=vision_hidden_size,
sd=self.sd_ref()
)
elif self.adapter_type == 'text_encoder':
@@ -731,7 +738,14 @@ class CustomAdapter(torch.nn.Module):
clip_image, output_hidden_states=True
)
img_embeds = id_embeds['last_hidden_state']
if self.config.clip_layer == 'penultimate_hidden_states':
img_embeds = id_embeds.hidden_states[-2]
elif self.config.clip_layer == 'last_hidden_state':
img_embeds = id_embeds.hidden_states[-1]
elif self.config.clip_layer == 'image_embeds':
img_embeds = id_embeds.image_embeds
else:
raise ValueError(f"unknown clip layer: {self.config.clip_layer}")
if self.config.quad_image:
# get the outputs of the quat

3
toolkit/models/ilora.py
View File

@@ -106,6 +106,9 @@ class InstantLoRAModule(torch.nn.Module):
# this will be used to add the vector to the original forward
def forward(self, img_embeds):
# expand token rank if only rank 2
if len(img_embeds.shape) == 2:
img_embeds = img_embeds.unsqueeze(1)
img_embeds = self.resampler(img_embeds)
self.img_embeds = img_embeds

16
toolkit/stable_diffusion_model.py
View File

@@ -863,6 +863,7 @@ class StableDiffusion:
unconditional_embeddings: Union[PromptEmbeds, None] = None,
is_input_scaled=False,
detach_unconditional=False,
rescale_cfg=None,
**kwargs,
):
# get the embeddings
@@ -1111,6 +1112,21 @@ class StableDiffusion:
noise_pred = noise_pred_uncond + guidance_scale * (
noise_pred_text - noise_pred_uncond
)
if rescale_cfg is not None and rescale_cfg != guidance_scale:
with torch.no_grad():
# do cfg at the target rescale so we can match it
target_pred_mean_std = noise_pred_uncond + rescale_cfg * (
noise_pred_text - noise_pred_uncond
)
target_mean = target_pred_mean_std.mean([1, 2, 3], keepdim=True).detach()
target_std = target_pred_mean_std.std([1, 2, 3], keepdim=True).detach()
pred_mean = noise_pred.mean([1, 2, 3], keepdim=True).detach()
pred_std = noise_pred.std([1, 2, 3], keepdim=True).detach()
# match the mean and std
noise_pred = (noise_pred - pred_mean) / pred_std
noise_pred = (noise_pred * target_std) + target_mean
# https://github.com/huggingface/diffusers/blob/7a91ea6c2b53f94da930a61ed571364022b21044/src/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl.py#L775
if guidance_rescale > 0.0: