Update README.md

README.md

@@ -153,10 +153,16 @@ def encode_prompt(
         input_ids, attention_mask=None, output_hidden_states=True
     )
 
-    start_embed = text_encoder_output.hidden_states[-1][:,0].unsqueeze(0)
-    end_embed = text_encoder_output.hidden_states[-1][:,-1].unsqueeze(0)
+    start_embed = text_encoder_output.hidden_states[-1][0][0].unsqueeze(0).unsqueeze(0)
+    end_embed = text_encoder_output.hidden_states[-1][0][0].unsqueeze(0).unsqueeze(0)
     prompt_embeds = text_encoder_output.hidden_states[-1][:,1:-1].reshape(1,-1,1280)
-    prompt_embeds = torch.cat([start_embed, prompt_embeds, end_embed], dim=1)
+
+    padding = []
+    for i in range((max_len + 1) - (prompt_embeds.shape[1] % 77)):
+        padding.append(end_embed)
+    padding = torch.cat(padding, dim=1)
+
+    prompt_embeds = torch.cat([start_embed, prompt_embeds, padding], dim=1)
     prompt_embeds = prompt_embeds.to(dtype=prior_pipe.text_encoder.dtype, device=device)
     prompt_embeds = prompt_embeds.repeat_interleave(num_images_per_prompt, dim=0)