File size: 2,864 Bytes
0305a63 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 |
import torch
def tokenize_prompt(tokenizer, prompt):
text_inputs = tokenizer(
prompt,
padding="max_length",
max_length=tokenizer.model_max_length,
truncation=True,
return_tensors="pt",
)
text_input_ids = text_inputs.input_ids
return text_input_ids
# Adapted from pipelines.StableDiffusionXLPipeline.encode_prompt
def encode_prompt(text_encoders, tokenizers, prompt, text_input_ids_list=None):
prompt_embeds_list = []
for i, text_encoder in enumerate(text_encoders):
if tokenizers is not None:
tokenizer = tokenizers[i]
text_input_ids = tokenize_prompt(tokenizer, prompt)
else:
assert text_input_ids_list is not None
text_input_ids = text_input_ids_list[i]
prompt_embeds = text_encoder(
text_input_ids.to(text_encoder.device),
output_hidden_states=True,
)
# We are only ALWAYS interested in the pooled output of the final text encoder
pooled_prompt_embeds = prompt_embeds[0]
prompt_embeds = prompt_embeds.hidden_states[-2]
bs_embed, seq_len, _ = prompt_embeds.shape
prompt_embeds = prompt_embeds.view(bs_embed, seq_len, -1)
prompt_embeds_list.append(prompt_embeds)
prompt_embeds = torch.concat(prompt_embeds_list, dim=-1)
pooled_prompt_embeds = pooled_prompt_embeds.view(bs_embed, -1)
return prompt_embeds, pooled_prompt_embeds
def add_tokens(tokenizers, tokens, text_encoders):
new_token_indices = {}
for idx, tokenizer in enumerate(tokenizers):
for token in tokens:
num_added_tokens = tokenizer.add_tokens(token)
if num_added_tokens == 0:
raise ValueError(
f"The tokenizer already contains the token {token}. Please pass a different"
" `placeholder_token` that is not already in the tokenizer."
)
new_token_indices[f"{idx}_{token}"] = num_added_tokens
# resize embedding layers to avoid crash. We will never actually use these.
text_encoders[idx].resize_token_embeddings(len(tokenizer), pad_to_multiple_of=128)
return new_token_indices
def patch_embedding_forward(embedding_layer, new_tokens, new_embeddings):
def new_forward(input):
embedded_text = torch.nn.functional.embedding(
input, embedding_layer.weight, embedding_layer.padding_idx, embedding_layer.max_norm,
embedding_layer.norm_type, embedding_layer.scale_grad_by_freq, embedding_layer.sparse)
replace_indices = (input == new_tokens)
if torch.count_nonzero(replace_indices) > 0:
embedded_text[replace_indices] = new_embeddings
return embedded_text
embedding_layer.forward = new_forward |