From d69cc8f46f238e91e2f597cd301cc53b1d4b8bec Mon Sep 17 00:00:00 2001 From: Volpeon Date: Sat, 25 Mar 2023 17:49:30 +0100 Subject: Fix training with guidance --- data/csv.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) (limited to 'data') diff --git a/data/csv.py b/data/csv.py index a6cd065..d52d251 100644 --- a/data/csv.py +++ b/data/csv.py @@ -104,11 +104,10 @@ def collate_fn(dtype: torch.dtype, tokenizer: CLIPTokenizer, with_guidance: bool nprompt_ids = [example["nprompt_ids"] for example in examples] input_ids = [example["instance_prompt_ids"] for example in examples] + negative_input_ids = [example["negative_prompt_ids"] for example in examples] pixel_values = [example["instance_images"] for example in examples] - if with_guidance: - input_ids += [example["negative_prompt_ids"] for example in examples] - elif with_prior_preservation: + if with_prior_preservation: input_ids += [example["class_prompt_ids"] for example in examples] pixel_values += [example["class_images"] for example in examples] @@ -118,13 +117,16 @@ def collate_fn(dtype: torch.dtype, tokenizer: CLIPTokenizer, with_guidance: bool prompts = unify_input_ids(tokenizer, prompt_ids) nprompts = unify_input_ids(tokenizer, nprompt_ids) inputs = unify_input_ids(tokenizer, input_ids) + negative_inputs = unify_input_ids(tokenizer, negative_input_ids) batch = { "prompt_ids": prompts.input_ids, "nprompt_ids": nprompts.input_ids, "input_ids": inputs.input_ids, + "negative_input_ids": negative_inputs.attention_mask, "pixel_values": pixel_values, "attention_mask": inputs.attention_mask, + "negative_attention_mask": negative_inputs.attention_mask, } return batch -- cgit v1.2.3-70-g09d2