Update README.md

README.md

@@ -153,16 +153,7 @@ def encode_prompt(
         input_ids, attention_mask=None, output_hidden_states=True
     )
 
-    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)
-
-    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 = text_encoder_output.hidden_states[-1].reshape(1,-1,1280)
     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)