| 92 | |
| 93 | |
| 94 | def encode_prompt(self, prompt_batch, proportion_empty_prompts, is_train=True): |
| 95 | prompt_embeds_list = [] |
| 96 | |
| 97 | captions = [] |
| 98 | for caption in prompt_batch: |
| 99 | if random.random() < proportion_empty_prompts: |
| 100 | captions.append("") |
| 101 | elif isinstance(caption, str): |
| 102 | captions.append(caption) |
| 103 | elif isinstance(caption, (list, np.ndarray)): |
| 104 | # take a random caption if there are multiple |
| 105 | captions.append(random.choice(caption) if is_train else caption[0]) |
| 106 | |
| 107 | with torch.no_grad(): |
| 108 | for tokenizer, text_encoder in zip(self.tokenizers, self.text_encoders): |
| 109 | text_inputs = tokenizer( |
| 110 | captions, |
| 111 | padding="max_length", |
| 112 | max_length=tokenizer.model_max_length, |
| 113 | truncation=True, |
| 114 | return_tensors="pt", |
| 115 | ) |
| 116 | text_input_ids = text_inputs.input_ids |
| 117 | prompt_embeds = text_encoder( |
| 118 | text_input_ids.to(text_encoder.device), |
| 119 | output_hidden_states=True, |
| 120 | ) |
| 121 | |
| 122 | # We are only ALWAYS interested in the pooled output of the final text encoder |
| 123 | pooled_prompt_embeds = prompt_embeds[0] |
| 124 | prompt_embeds = prompt_embeds.hidden_states[-2] |
| 125 | bs_embed, seq_len, _ = prompt_embeds.shape |
| 126 | prompt_embeds = prompt_embeds.view(bs_embed, seq_len, -1) |
| 127 | prompt_embeds_list.append(prompt_embeds) |
| 128 | |
| 129 | prompt_embeds = torch.concat(prompt_embeds_list, dim=-1) |
| 130 | pooled_prompt_embeds = pooled_prompt_embeds.view(bs_embed, -1) |
| 131 | return prompt_embeds, pooled_prompt_embeds |
| 132 | |
| 133 | def compute_embeddings(self, prompt_batch, proportion_empty_prompts, text_encoders, tokenizers, size, is_train=True): |
| 134 | original_size = size |