Fix training with guidance

1 files changed, 5 insertions, 3 deletions

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