2023-02-22 05:50:27 +00:00
|
|
|
from logging import getLogger
|
2023-02-22 04:50:59 +00:00
|
|
|
from os import makedirs, path
|
2023-03-15 22:14:52 +00:00
|
|
|
from typing import List, Optional, Tuple
|
2023-02-22 05:50:27 +00:00
|
|
|
|
2023-03-15 22:14:52 +00:00
|
|
|
import numpy as np
|
2023-02-22 05:50:27 +00:00
|
|
|
import torch
|
2023-02-21 05:07:16 +00:00
|
|
|
from huggingface_hub.file_download import hf_hub_download
|
2023-03-15 22:14:52 +00:00
|
|
|
from onnx import ModelProto, load_model, numpy_helper, save_model
|
|
|
|
from transformers import CLIPTokenizer
|
2023-02-21 05:07:16 +00:00
|
|
|
|
2023-03-15 22:14:52 +00:00
|
|
|
from ...server.context import ServerContext
|
2023-02-22 03:40:57 +00:00
|
|
|
from ..utils import ConversionContext
|
2023-02-21 05:07:16 +00:00
|
|
|
|
|
|
|
logger = getLogger(__name__)
|
|
|
|
|
|
|
|
|
2023-03-15 22:14:52 +00:00
|
|
|
@torch.no_grad()
|
|
|
|
def blend_textual_inversions(
|
|
|
|
context: ServerContext,
|
|
|
|
text_encoder: Optional[ModelProto],
|
|
|
|
tokenizer: Optional[CLIPTokenizer],
|
|
|
|
inversion_names: List[str],
|
|
|
|
inversion_formats: List[str],
|
|
|
|
inversion_weights: Optional[List[float]] = None,
|
|
|
|
base_tokens: Optional[List[str]] = None,
|
|
|
|
) -> Tuple[ModelProto, CLIPTokenizer]:
|
2023-03-16 00:27:29 +00:00
|
|
|
dtype = np.float
|
2023-03-15 22:14:52 +00:00
|
|
|
embeds = {}
|
|
|
|
|
2023-03-16 00:27:29 +00:00
|
|
|
for name, format, weight, base_token in zip(
|
|
|
|
inversion_names,
|
|
|
|
inversion_formats,
|
|
|
|
inversion_weights,
|
|
|
|
base_tokens or inversion_names,
|
|
|
|
):
|
2023-03-15 22:14:52 +00:00
|
|
|
logger.info("blending Textual Inversion %s with weight of %s", name, weight)
|
|
|
|
if format == "concept":
|
|
|
|
embeds_file = hf_hub_download(repo_id=name, filename="learned_embeds.bin")
|
|
|
|
token_file = hf_hub_download(repo_id=name, filename="token_identifier.txt")
|
|
|
|
|
|
|
|
with open(token_file, "r") as f:
|
|
|
|
token = base_token or f.read()
|
|
|
|
|
|
|
|
loaded_embeds = torch.load(embeds_file)
|
|
|
|
|
|
|
|
# separate token and the embeds
|
|
|
|
trained_token = list(loaded_embeds.keys())[0]
|
|
|
|
|
|
|
|
layer = loaded_embeds[trained_token].cpu().numpy().astype(dtype)
|
|
|
|
layer *= weight
|
|
|
|
if trained_token in embeds:
|
|
|
|
embeds[token] += layer
|
|
|
|
else:
|
|
|
|
embeds[token] = layer
|
|
|
|
elif format == "embeddings":
|
|
|
|
loaded_embeds = torch.load(name)
|
|
|
|
|
|
|
|
string_to_token = loaded_embeds["string_to_token"]
|
|
|
|
string_to_param = loaded_embeds["string_to_param"]
|
|
|
|
|
|
|
|
# separate token and embeds
|
|
|
|
trained_token = list(string_to_token.keys())[0]
|
|
|
|
trained_embeds = string_to_param[trained_token]
|
|
|
|
|
|
|
|
num_tokens = trained_embeds.shape[0]
|
2023-03-17 00:37:25 +00:00
|
|
|
logger.debug("generating %s layer tokens for %s", num_tokens, name)
|
2023-03-15 22:14:52 +00:00
|
|
|
|
|
|
|
for i in range(num_tokens):
|
|
|
|
token = f"{base_token or name}-{i}"
|
2023-03-16 00:27:29 +00:00
|
|
|
layer = trained_embeds[i, :].cpu().numpy().astype(dtype)
|
2023-03-15 22:14:52 +00:00
|
|
|
layer *= weight
|
|
|
|
if token in embeds:
|
|
|
|
embeds[token] += layer
|
|
|
|
else:
|
|
|
|
embeds[token] = layer
|
|
|
|
else:
|
|
|
|
raise ValueError(f"unknown Textual Inversion format: {format}")
|
|
|
|
|
|
|
|
# add the tokens to the tokenizer
|
2023-03-17 00:37:25 +00:00
|
|
|
logger.debug(
|
2023-03-16 00:27:29 +00:00
|
|
|
"found embeddings for %s tokens: %s", len(embeds.keys()), embeds.keys()
|
|
|
|
)
|
2023-03-15 22:14:52 +00:00
|
|
|
num_added_tokens = tokenizer.add_tokens(list(embeds.keys()))
|
|
|
|
if num_added_tokens == 0:
|
|
|
|
raise ValueError(
|
|
|
|
f"The tokenizer already contains the token {token}. Please pass a different `token` that is not already in the tokenizer."
|
|
|
|
)
|
|
|
|
|
2023-03-17 00:37:25 +00:00
|
|
|
logger.trace("added %s tokens", num_added_tokens)
|
2023-03-15 22:14:52 +00:00
|
|
|
|
|
|
|
# resize the token embeddings
|
|
|
|
# text_encoder.resize_token_embeddings(len(tokenizer))
|
2023-03-16 00:27:29 +00:00
|
|
|
embedding_node = [
|
|
|
|
n
|
|
|
|
for n in text_encoder.graph.initializer
|
|
|
|
if n.name == "text_model.embeddings.token_embedding.weight"
|
|
|
|
][0]
|
2023-03-15 22:14:52 +00:00
|
|
|
embedding_weights = numpy_helper.to_array(embedding_node)
|
|
|
|
|
|
|
|
weights_dim = embedding_weights.shape[1]
|
|
|
|
zero_weights = np.zeros((num_added_tokens, weights_dim))
|
|
|
|
embedding_weights = np.concatenate((embedding_weights, zero_weights), axis=0)
|
|
|
|
|
|
|
|
for token, weights in embeds.items():
|
|
|
|
token_id = tokenizer.convert_tokens_to_ids(token)
|
2023-03-17 00:37:25 +00:00
|
|
|
logger.trace("embedding %s weights for token %s", weights.shape, token)
|
2023-03-15 22:14:52 +00:00
|
|
|
embedding_weights[token_id] = weights
|
|
|
|
|
|
|
|
# replace embedding_node
|
|
|
|
for i in range(len(text_encoder.graph.initializer)):
|
2023-03-16 00:27:29 +00:00
|
|
|
if (
|
|
|
|
text_encoder.graph.initializer[i].name
|
|
|
|
== "text_model.embeddings.token_embedding.weight"
|
|
|
|
):
|
|
|
|
new_initializer = numpy_helper.from_array(
|
|
|
|
embedding_weights.astype(np.float32), embedding_node.name
|
|
|
|
)
|
2023-03-17 00:37:25 +00:00
|
|
|
logger.trace("new initializer data type: %s", new_initializer.data_type)
|
2023-03-15 22:14:52 +00:00
|
|
|
del text_encoder.graph.initializer[i]
|
|
|
|
text_encoder.graph.initializer.insert(i, new_initializer)
|
|
|
|
|
|
|
|
return (text_encoder, tokenizer)
|
|
|
|
|
|
|
|
|
2023-03-01 14:26:40 +00:00
|
|
|
@torch.no_grad()
|
2023-02-22 05:50:27 +00:00
|
|
|
def convert_diffusion_textual_inversion(
|
2023-03-05 13:19:48 +00:00
|
|
|
context: ConversionContext,
|
|
|
|
name: str,
|
|
|
|
base_model: str,
|
|
|
|
inversion: str,
|
|
|
|
format: str,
|
|
|
|
base_token: Optional[str] = None,
|
2023-02-22 05:50:27 +00:00
|
|
|
):
|
2023-02-22 04:49:34 +00:00
|
|
|
dest_path = path.join(context.model_path, f"inversion-{name}")
|
2023-02-22 05:50:27 +00:00
|
|
|
logger.info(
|
|
|
|
"converting Textual Inversion: %s + %s -> %s", base_model, inversion, dest_path
|
|
|
|
)
|
2023-02-21 05:07:16 +00:00
|
|
|
|
2023-03-08 05:40:04 +00:00
|
|
|
encoder_path = path.join(dest_path, "text_encoder")
|
|
|
|
encoder_model = path.join(encoder_path, "model.onnx")
|
|
|
|
tokenizer_path = path.join(dest_path, "tokenizer")
|
|
|
|
|
2023-03-08 05:57:39 +00:00
|
|
|
if (
|
|
|
|
path.exists(dest_path)
|
|
|
|
and path.exists(encoder_model)
|
|
|
|
and path.exists(tokenizer_path)
|
|
|
|
):
|
2023-02-22 03:40:57 +00:00
|
|
|
logger.info("ONNX model already exists, skipping.")
|
2023-02-22 05:50:27 +00:00
|
|
|
return
|
2023-02-22 03:40:57 +00:00
|
|
|
|
2023-03-08 05:40:04 +00:00
|
|
|
makedirs(encoder_path, exist_ok=True)
|
2023-02-21 05:07:16 +00:00
|
|
|
|
2023-03-15 22:14:52 +00:00
|
|
|
text_encoder = load_model(path.join(base_model, "text_encoder", "model.onnx"))
|
2023-02-21 05:07:16 +00:00
|
|
|
tokenizer = CLIPTokenizer.from_pretrained(
|
|
|
|
base_model,
|
|
|
|
subfolder="tokenizer",
|
2023-02-22 04:49:34 +00:00
|
|
|
)
|
2023-03-15 22:14:52 +00:00
|
|
|
text_encoder, tokenizer = blend_textual_inversions(
|
|
|
|
context,
|
|
|
|
text_encoder,
|
|
|
|
tokenizer,
|
|
|
|
[inversion],
|
|
|
|
[format],
|
|
|
|
[1.0],
|
|
|
|
base_token=(base_token or name),
|
2023-02-21 05:07:16 +00:00
|
|
|
)
|
|
|
|
|
2023-03-02 13:57:59 +00:00
|
|
|
logger.info("saving tokenizer for textual inversion")
|
2023-03-08 05:40:04 +00:00
|
|
|
tokenizer.save_pretrained(tokenizer_path)
|
2023-03-02 01:08:31 +00:00
|
|
|
|
2023-03-02 13:57:59 +00:00
|
|
|
logger.info("saving text encoder for textual inversion")
|
2023-03-15 22:14:52 +00:00
|
|
|
save_model(
|
2023-02-21 05:07:16 +00:00
|
|
|
text_encoder,
|
2023-03-08 05:40:04 +00:00
|
|
|
f=encoder_model,
|
2023-02-21 05:07:16 +00:00
|
|
|
)
|
2023-03-02 13:57:59 +00:00
|
|
|
|
|
|
|
logger.info("textual inversion saved to %s", dest_path)
|