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

added prompt dropout to happen indempendently on each TE

Browse Source
This commit is contained in:
Jaret Burkett
2023-11-14 05:26:51 -07:00
parent 7782caa468
commit 4f9cdd916a
7 changed files with 144 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
toolkit/config_modules.py
View File

@@ -194,6 +194,9 @@ class TrainConfig:
self.noise_multiplier = kwargs.get('noise_multiplier', 1.0)
self.img_multiplier = kwargs.get('img_multiplier', 1.0)
# dropout that happens before encoding. It functions independently per text encoder
self.prompt_dropout_prob = kwargs.get('prompt_dropout_prob', 0.0)
# match the norm of the noise before computing loss. This will help the model maintain its
# current understandin of the brightness of images.

5
toolkit/sampler.py
View File

@@ -12,7 +12,9 @@ from diffusers import (
HeunDiscreteScheduler,
KDPM2DiscreteScheduler,
KDPM2AncestralDiscreteScheduler,
LCMScheduler
)
from k_diffusion.external import CompVisDenoiser
# scheduler:
@@ -72,12 +74,15 @@ def get_sampler(
scheduler_cls = KDPM2DiscreteScheduler
elif sampler == "dpm_2_a":
scheduler_cls = KDPM2AncestralDiscreteScheduler
elif sampler == "lcm":
scheduler_cls = LCMScheduler
config = copy.deepcopy(sdxl_sampler_config)
config.update(sched_init_args)
scheduler = scheduler_cls.from_config(config)
return scheduler

30
toolkit/stable_diffusion_model.py
View File

@@ -344,6 +344,11 @@ class StableDiffusion:
else:
noise_scheduler = get_sampler(sampler)
try:
noise_scheduler = noise_scheduler.to(self.device_torch, self.torch_dtype)
except:
pass
if sampler.startswith("sample_") and self.is_xl:
# using kdiffusion
Pipe = StableDiffusionKDiffusionXLPipeline
@@ -722,7 +727,8 @@ class StableDiffusion:
refiner_pred = self.refiner_unet(
input_chunks[1],
timestep_chunks[1],
encoder_hidden_states=text_embeds_chunks[1][:, :, -1280:], # just use the first second text encoder
encoder_hidden_states=text_embeds_chunks[1][:, :, -1280:],
# just use the first second text encoder
added_cond_kwargs={
"text_embeds": added_cond_kwargs_chunked['text_embeds'][1],
# "time_ids": added_cond_kwargs_chunked['time_ids'][1],
@@ -740,7 +746,8 @@ class StableDiffusion:
# just use the first second text encoder
added_cond_kwargs={
"text_embeds": text_embeddings.pooled_embeds,
"time_ids": self.get_time_ids_from_latents(latent_model_input, requires_aesthetic_score=True),
"time_ids": self.get_time_ids_from_latents(latent_model_input,
requires_aesthetic_score=True),
},
**kwargs,
).sample
@@ -845,7 +852,8 @@ class StableDiffusion:
num_images_per_prompt=1,
force_all=False,
long_prompts=False,
max_length=None
max_length=None,
dropout_prob=0.0,
) -> PromptEmbeds:
# sd1.5 embeddings are (bs, 77, 768)
prompt = prompt
@@ -875,12 +883,18 @@ class StableDiffusion:
use_text_encoder_2=use_encoder_2,
truncate=not long_prompts,
max_length=max_length,
dropout_prob=dropout_prob,
)
)
else:
return PromptEmbeds(
train_tools.encode_prompts(
self.tokenizer, self.text_encoder, prompt, truncate=not long_prompts, max_length=max_length
self.tokenizer,
self.text_encoder,
prompt,
truncate=not long_prompts,
max_length=max_length,
dropout_prob=dropout_prob
)
)
@@ -1011,8 +1025,9 @@ class StableDiffusion:
state_dict[new_key] = v
return state_dict
def named_parameters(self, vae=True, text_encoder=True, unet=True, refiner=False, state_dict_keys=False) -> OrderedDict[
str, Parameter]:
def named_parameters(self, vae=True, text_encoder=True, unet=True, refiner=False, state_dict_keys=False) -> \
OrderedDict[
str, Parameter]:
named_params: OrderedDict[str, Parameter] = OrderedDict()
if vae:
for name, param in self.vae.named_parameters(recurse=True, prefix=f"{SD_PREFIX_VAE}"):
@@ -1198,7 +1213,8 @@ class StableDiffusion:
print(f"Found {len(params)} trainable parameter in text encoder")
if refiner:
named_params = self.named_parameters(vae=False, unet=False, text_encoder=False, refiner=True, state_dict_keys=True)
named_params = self.named_parameters(vae=False, unet=False, text_encoder=False, refiner=True,
state_dict_keys=True)
refiner_lr = refiner_lr if refiner_lr is not None else default_lr
params = []
for key, diffusers_key in ldm_diffusers_keymap.items():

15
toolkit/train_tools.py
View File

@@ -537,6 +537,7 @@ def encode_prompts_xl(
use_text_encoder_2: bool = True, # sdxl
truncate: bool = True,
max_length=None,
dropout_prob=0.0,
) -> tuple[torch.FloatTensor, torch.FloatTensor]:
# text_encoder and text_encoder_2's penuultimate layer's output
text_embeds_list = []
@@ -553,6 +554,12 @@ def encode_prompts_xl(
if idx == 1 and not use_text_encoder_2:
prompt_list_to_use = ["" for _ in prompts]
if dropout_prob > 0.0:
# randomly drop out prompts
prompt_list_to_use = [
prompt if torch.rand(1).item() > dropout_prob else "" for prompt in prompt_list_to_use
]
text_tokens_input_ids = text_tokenize(tokenizer, prompt_list_to_use, truncate=truncate, max_length=max_length)
# set the max length for the next one
if idx == 0:
@@ -598,9 +605,17 @@ def encode_prompts(
prompts: list[str],
truncate: bool = True,
max_length=None,
dropout_prob=0.0,
):
if max_length is None:
max_length = tokenizer.model_max_length
if dropout_prob > 0.0:
# randomly drop out prompts
prompts = [
prompt if torch.rand(1).item() > dropout_prob else "" for prompt in prompts
]
text_tokens = text_tokenize(tokenizer, prompts, truncate=truncate, max_length=max_length)
text_embeddings = text_encode(text_encoder, text_tokens, truncate=truncate, max_length=max_length)