mirror of
https://github.com/modelscope/DiffSynth-Studio.git
synced 2026-03-19 06:48:12 +00:00
73 lines
3.4 KiB
Python
73 lines
3.4 KiB
Python
from .base_prompter import BasePrompter, tokenize_long_prompt
|
|
from ..models.model_manager import ModelManager, 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)
|
|
|
|
# Search for embeddings
|
|
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 |