Update README.md

README.md

@@ -124,33 +124,45 @@ def encode_prompt(
     prompt=""
     ):
 
-    text_inputs = prior_pipe.tokenizer(
-        prompt,
-        padding="longest",
-        truncation=False,
-        return_tensors="pt",
-    )
-    chunk = []
-    padding = []
-    max_len = 75
-    start_token = text_inputs.input_ids[:,0].unsqueeze(0)
-    end_token = text_inputs.input_ids[:,-1].unsqueeze(0)
-    raw_input_ids = text_inputs.input_ids[:,1:-1]
-    prompt_len = len(raw_input_ids[0])
-    last_lenght = prompt_len % max_len
-    
-    for i in range(int((prompt_len - last_lenght) / max_len)):
-        chunk.append(torch.cat([start_token, raw_input_ids[:,i*max_len:(i+1)*max_len], end_token], dim=1))
-    for i in range(max_len - last_lenght):
-        padding.append(text_inputs.input_ids[:,-1])
-    
-    last_chunk = torch.cat([raw_input_ids[:,prompt_len-last_lenght:], torch.tensor([padding])], dim=1)
-    chunk.append(torch.cat([start_token, last_chunk, end_token], dim=1))
-    input_ids = torch.cat(chunk, dim=0).to(device)
-
-    # Don't use attention masks
+    if prompt == "":
+        text_inputs = prior_pipe.tokenizer(
+            prompt,
+            padding="max_length",
+            max_length=77,
+            truncation=False,
+            return_tensors="pt",
+        )
+        input_ids = text_inputs.input_ids
+        attention_mask=None
+    else:   
+        text_inputs = prior_pipe.tokenizer(
+            prompt,
+            padding="longest",
+            truncation=False,
+            return_tensors="pt",
+        )
+        chunk = []
+        padding = []
+        max_len = 75
+        start_token = text_inputs.input_ids[:,0].unsqueeze(0)
+        end_token = text_inputs.input_ids[:,-1].unsqueeze(0)
+        raw_input_ids = text_inputs.input_ids[:,1:-1]
+        prompt_len = len(raw_input_ids[0])
+        last_lenght = prompt_len % max_len
+        
+        for i in range(int((prompt_len - last_lenght) / max_len)):
+            chunk.append(torch.cat([start_token, raw_input_ids[:,i*max_len:(i+1)*max_len], end_token], dim=1))
+        for i in range(max_len - last_lenght):
+            padding.append(text_inputs.input_ids[:,-1])
+        
+        last_chunk = torch.cat([raw_input_ids[:,prompt_len-last_lenght:], torch.tensor([padding])], dim=1)
+        chunk.append(torch.cat([start_token, last_chunk, end_token], dim=1))
+        input_ids = torch.cat(chunk, dim=0)
+        attention_mask = torch.ones(input_ids.shape, device=device, dtype=torch.int64)
+        attention_mask[-1,last_lenght+1:] = 0
+
     text_encoder_output = prior_pipe.text_encoder(
-        input_ids, attention_mask=None, output_hidden_states=True
+        input_ids.to(device), attention_mask=attention_mask, output_hidden_states=True
     )
 
     prompt_embeds = text_encoder_output.hidden_states[-1].reshape(1,-1,1280)
@@ -169,7 +181,7 @@ quality_prompt = "very aesthetic, best quality, newest"
 negative_prompt = "very displeasing, displeasing, worst quality, bad quality, low quality, realistic, monochrome, comic, sketch, oldest, early, artist name, signature, blurry, simple background, upside down, interlocked fingers,"
 num_images_per_prompt=1
 
-# Encode prompts and quality prompts eperately, don't use attention masks and long prompt support:
+# Encode prompts and quality prompts eperately, long prompt support and don't use attention masks for empty prompts:
 # pipe, device, num_images_per_prompt, prompt
 empty_prompt_embeds, _ = encode_prompt(pipe.prior_pipe, device, num_images_per_prompt, prompt="")