Update README.md

README.md

@@ -117,29 +117,54 @@ pipe = pipe.to(device, dtype=dtype)
 pipe.prior_pipe = pipe.prior_pipe.to(device, dtype=dtype)
 
 
-def encode_empty_prompt(
+def encode_prompt(
     prior_pipe,
     device,
-    batch_size,
     num_images_per_prompt,
+    prompt=""
     ):
 
     text_inputs = prior_pipe.tokenizer(
-        "",
-        padding="max_length",
-        max_length=prior_pipe.tokenizer.model_max_length,
-        truncation=True,
+        prompt,
+        padding="longest",
+        truncation=False,
         return_tensors="pt",
     )
-
-    # Don't use attention mask for empty prompt
+    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
     text_encoder_output = prior_pipe.text_encoder(
-        text_inputs.input_ids.to(device), attention_mask=None, output_hidden_states=True
+        input_ids, attention_mask=None, output_hidden_states=True
     )
-    prompt_embeds = text_encoder_output.hidden_states[-1]
+
+    start_embed = text_encoder_output.hidden_states[-1][:,0].unsqueeze(0)
+    end_embed = text_encoder_output.hidden_states[-1][:,-1].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)
     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)
-    return prompt_embeds
+
+    prompt_embeds_pooled = text_encoder_output.text_embeds[0].unsqueeze(0).unsqueeze(1)
+    prompt_embeds_pooled = prompt_embeds_pooled.to(dtype=prior_pipe.text_encoder.dtype, device=device)
+    prompt_embeds_pooled = prompt_embeds_pooled.repeat_interleave(num_images_per_prompt, dim=0)
+
+    return prompt_embeds, prompt_embeds_pooled
 
 
 prompt = "1girl, solo, looking at viewer, open mouth, blue eyes, medium breasts, blonde hair, gloves, dress, bow, hair between eyes, bare shoulders, upper body, hair bow, indoors, elbow gloves, hand on own chest, bridal gauntlets, candlestand, smile, rim lighting, from side, castle interior, looking side,"
@@ -147,18 +172,21 @@ quality_prompt = "extremely 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 seperately:
-# device, batch_size, num_images_per_prompt, cfg, prompt
-prompt_embeds, prompt_embeds_pooled, _, _ = pipe.prior_pipe.encode_prompt(device, 1, num_images_per_prompt, False, prompt=prompt)
-quality_prompt_embeds, _, _, _ = pipe.prior_pipe.encode_prompt(device, 1, num_images_per_prompt, False, prompt=quality_prompt)
-
-negative_prompt_embeds, negative_prompt_embeds_pooled, _, _ = pipe.prior_pipe.encode_prompt(device, 1, num_images_per_prompt, False, prompt=negative_prompt)
-empty_prompt_embeds = encode_empty_prompt(pipe.prior_pipe, device, 1, num_images_per_prompt)
+# Encode prompts and quality prompts eperately, don't use attention masks and long prompt support:
+# pipe, device, num_images_per_prompt, prompt
+empty_prompt_embeds, _ = encode_prompt(pipe.prior_pipe, device, num_images_per_prompt, prompt="")
 
+prompt_embeds, prompt_embeds_pooled = encode_prompt(pipe.prior_pipe, device, num_images_per_prompt, prompt=prompt)
+quality_prompt_embeds, _ = encode_prompt(pipe.prior_pipe, device, num_images_per_prompt, prompt=quality_prompt)
 prompt_embeds = torch.cat([prompt_embeds, quality_prompt_embeds], dim=1)
-negative_prompt_embeds = torch.cat([negative_prompt_embeds, empty_prompt_embeds], dim=1)
 
-pipe.prior_pipe.maybe_free_model_hooks()
+negative_prompt_embeds, negative_prompt_embeds_pooled = encode_prompt(pipe.prior_pipe, device, num_images_per_prompt, prompt=negative_prompt)
+
+while prompt_embeds.shape[1] < negative_prompt_embeds.shape[1]:
+    prompt_embeds = torch.cat([prompt_embeds, empty_prompt_embeds], dim=1)
+
+while negative_prompt_embeds.shape[1] < prompt_embeds.shape[1]:
+    negative_prompt_embeds = torch.cat([negative_prompt_embeds, empty_prompt_embeds], dim=1)
 
 output = pipe(
     width=1024,