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

Added target_norm_std which is a game changer

Browse Source
This commit is contained in:
Jaret Burkett
2024-07-28 16:08:33 -06:00
parent 0bc4d555c7
commit e81e19fd0f
4 changed files with 11 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
toolkit/train_tools.py
View File

@@ -751,7 +751,6 @@ def encode_prompts_auraflow(
padding="max_length",
return_tensors="pt",
)
text_inputs = {k: v.to(device) for k, v in text_inputs.items()}
text_input_ids = text_inputs["input_ids"]
untruncated_ids = tokenizer(prompts, padding="longest", return_tensors="pt").input_ids
@@ -760,6 +759,7 @@ def encode_prompts_auraflow(
):
removed_text = tokenizer.batch_decode(untruncated_ids[:, max_length - 1: -1])
text_inputs = {k: v.to(device) for k, v in text_inputs.items()}
prompt_embeds = text_encoder(**text_inputs)[0]
prompt_attention_mask = text_inputs["attention_mask"].unsqueeze(-1).expand(prompt_embeds.shape)
prompt_embeds = prompt_embeds * prompt_attention_mask