| | from .base_prompter import BasePrompter, tokenize_long_prompt |
| | from ..models.utils import load_state_dict, search_for_embeddings |
| | from ..models import SDTextEncoder |
| | from transformers import CLIPTokenizer |
| | import torch, os |
| |
|
| |
|
| |
|
| | class SDPrompter(BasePrompter): |
| | def __init__(self, tokenizer_path=None): |
| | if tokenizer_path is None: |
| | base_path = os.path.dirname(os.path.dirname(__file__)) |
| | tokenizer_path = os.path.join(base_path, "tokenizer_configs/stable_diffusion/tokenizer") |
| | super().__init__() |
| | self.tokenizer = CLIPTokenizer.from_pretrained(tokenizer_path) |
| | self.text_encoder: SDTextEncoder = None |
| | self.textual_inversion_dict = {} |
| | self.keyword_dict = {} |
| |
|
| |
|
| | def fetch_models(self, text_encoder: SDTextEncoder = None): |
| | self.text_encoder = text_encoder |
| |
|
| |
|
| | def add_textual_inversions_to_model(self, textual_inversion_dict, text_encoder): |
| | dtype = next(iter(text_encoder.parameters())).dtype |
| | state_dict = text_encoder.token_embedding.state_dict() |
| | token_embeddings = [state_dict["weight"]] |
| | for keyword in textual_inversion_dict: |
| | _, embeddings = textual_inversion_dict[keyword] |
| | token_embeddings.append(embeddings.to(dtype=dtype, device=token_embeddings[0].device)) |
| | token_embeddings = torch.concat(token_embeddings, dim=0) |
| | state_dict["weight"] = token_embeddings |
| | text_encoder.token_embedding = torch.nn.Embedding(token_embeddings.shape[0], token_embeddings.shape[1]) |
| | text_encoder.token_embedding = text_encoder.token_embedding.to(dtype=dtype, device=token_embeddings[0].device) |
| | text_encoder.token_embedding.load_state_dict(state_dict) |
| |
|
| |
|
| | def add_textual_inversions_to_tokenizer(self, textual_inversion_dict, tokenizer): |
| | additional_tokens = [] |
| | for keyword in textual_inversion_dict: |
| | tokens, _ = textual_inversion_dict[keyword] |
| | additional_tokens += tokens |
| | self.keyword_dict[keyword] = " " + " ".join(tokens) + " " |
| | tokenizer.add_tokens(additional_tokens) |
| |
|
| |
|
| | def load_textual_inversions(self, model_paths): |
| | for model_path in model_paths: |
| | keyword = os.path.splitext(os.path.split(model_path)[-1])[0] |
| | state_dict = load_state_dict(model_path) |
| |
|
| | |
| | for embeddings in search_for_embeddings(state_dict): |
| | if len(embeddings.shape) == 2 and embeddings.shape[1] == 768: |
| | tokens = [f"{keyword}_{i}" for i in range(embeddings.shape[0])] |
| | self.textual_inversion_dict[keyword] = (tokens, embeddings) |
| |
|
| | self.add_textual_inversions_to_model(self.textual_inversion_dict, self.text_encoder) |
| | self.add_textual_inversions_to_tokenizer(self.textual_inversion_dict, self.tokenizer) |
| |
|
| |
|
| | def encode_prompt(self, prompt, clip_skip=1, device="cuda", positive=True): |
| | prompt = self.process_prompt(prompt, positive=positive) |
| | for keyword in self.keyword_dict: |
| | if keyword in prompt: |
| | print(f"Textual inversion {keyword} is enabled.") |
| | prompt = prompt.replace(keyword, self.keyword_dict[keyword]) |
| | input_ids = tokenize_long_prompt(self.tokenizer, prompt).to(device) |
| | prompt_emb = self.text_encoder(input_ids, clip_skip=clip_skip) |
| | prompt_emb = prompt_emb.reshape((1, prompt_emb.shape[0]*prompt_emb.shape[1], -1)) |
| |
|
| | return prompt_emb |
