tokenizer_util.py

import torch
from typing import Callable, Dict, Optional, Union, List
from urllib.parse import urlparse
from transformers import PreTrainedModel, PreTrainedTokenizer, CLIPTokenizer
import copy
import random
from io import BytesIO
from compel.embeddings_provider import BaseTextualInversionManager

class TextualInversionLoaderMixin(BaseTextualInversionManager):
    r"""
    Mixin class for adding textual inversion tokens and embeddings to the tokenizer and text encoder with method:
    - [`~TextualInversionLoaderMixin.load_textual_inversion_embeddings`]
    - [`~TextualInversionLoaderMixin.add_textual_inversion_embedding`]
    """

    def load_textual_inversion_embeddings(
        self,
        embedding_path_dict_or_list: Union[Dict[str, str], List[Dict[str, str]]],
        allow_replacement: bool = False,
    ):
        r"""
        Loads textual inversion embeddings and adds them to the tokenizer's vocabulary and the text encoder's embeddings.
        Arguments:
            embeddings_path_dict_or_list (`Dict[str, str]` or `List[str]`):
                Dictionary of token to embedding path or List of embedding paths to embedding dictionaries.
                The dictionary must have the following keys:
                    - `token`: name of the token to be added to the tokenizers' vocabulary
                    - `embedding`: path to the embedding of the token to be added to the text encoder's embedding matrix
                The list must contain paths to embedding dictionaries where the keys are the tokens and the
                values are the embeddings (same as above dictionary definition).
            allow_replacement (`bool`, *optional*, defaults to `False`):
                Whether to allow replacement of existing tokens in the tokenizer's vocabulary. If `False`
                and a token is already in the vocabulary, an error will be raised.
        Returns:
            None
        """
        # Validate that inheriting class instance contains required attributes
        self._validate_method_call(self.load_textual_inversion_embeddings)

        if isinstance(embedding_path_dict_or_list, dict):
            for token, embedding_path in embedding_path_dict_or_list.items():
                
                embedding_dict = torch.load(embedding_path, map_location=self.text_encoder.device)
                embedding, is_multi_vec_token = self._extract_embedding_from_dict(embedding_dict)

                self._validate_token_update(token, allow_replacement, is_multi_vec_token)
                self.add_textual_inversion_embedding(token, embedding)
        elif isinstance(embedding_path_dict_or_list, list):
            for embedding_path in embedding_path_dict_or_list:
                embedding_dict = torch.load(embedding_path, map_location=self.text_encoder.device)
                token = self._extract_token_from_dict(embedding_dict)
                embedding, is_multi_vec_token = self._extract_embedding_from_dict(embedding_dict)

                self._validate_token_update(token, allow_replacement, is_multi_vec_token)
                self.add_textual_inversion_embedding(token, embedding)
        else:
            raise ValueError(
                f"Type {type(embedding_path_dict_or_list)} is invalid. The value passed to `embedding_path_dict_or_list` "
                "must be a dictionary that maps a token to it's embedding file path "
                "or a list of paths to embedding files containing embedding dictionaries."
            )

    def add_textual_inversion_embedding(self, token: str, embedding: torch.Tensor):
        r"""
        Adds a token to the tokenizer's vocabulary and an embedding to the text encoder's embedding matrix.
        Arguments:
            token (`str`):
                The token to be added to the tokenizers' vocabulary
            embedding (`torch.Tensor`):
                The embedding of the token to be added to the text encoder's embedding matrix
        Returns:
            None
        """
        # NOTE: Not clear to me that we intend for this to be a public/exposed method.
        # Validate that inheriting class instance contains required attributes
        self._validate_method_call(self.load_textual_inversion_embeddings)

        embedding = embedding.to(self.text_encoder.dtype)

        if not isinstance(self.tokenizer, MultiTokenCLIPTokenizer):
            if token in self.tokenizer.get_vocab():
                # If user has allowed replacement and the token exists, we only need to
                # extract the existing id and update the embedding
                token_id = self.tokenizer.convert_tokens_to_ids(token)
                self.text_encoder.get_input_embeddings().weight.data[token_id] = embedding
            else:
                # If the token does not exist, we add it to the tokenizer, then resize and update the
                # text encoder acccordingly
                self.tokenizer.add_tokens([token])

                token_id = self.tokenizer.convert_tokens_to_ids(token)
                self.text_encoder.resize_token_embeddings(len(self.tokenizer))
                self.text_encoder.get_input_embeddings().weight.data[token_id] = embedding
        else:
            if token in self.tokenizer.token_map:
                # If user has allowed replacement and the token exists, we need to
                # remove all existing tokens associated with the old embbedding and
                # upddate with the new ones
                indices_to_remove = []
                for token_to_remove in self.tokenizer.token_map[token]:
                    indices_to_remove.append(self.tokenizer.get_added_vocab()[token_to_remove])

                    # Remove old  tokens from tokenizer
                    self.tokenizer.added_tokens_encoder.pop(token_to_remove)

                # Convert indices to remove to tensor
                indices_to_remove = torch.LongTensor(indices_to_remove)

                # Remove old tokens from text encoder
                token_embeds = self.text_encoder.get_input_embeddings().weight.data
                indices_to_keep = torch.arange(0, token_embeds.shape[0])
                indices_to_keep = indices_to_keep[indices_to_keep != indices_to_remove].squeeze()
                token_embeds = token_embeds[indices_to_keep]

                # Downsize text encoder
                self.text_encoder.resize_token_embeddings(len(self.tokenizer))

                # Remove token from map so MultiTokenCLIPTokenizer doesn't complain
                # on update
                self.tokenizer.token_map.pop(token)

            # Update token with new embedding
            embedding_dims = len(embedding.shape)
            num_vec_per_token = 1 if embedding_dims == 1 else embedding.shape[0]

            self.tokenizer.add_placeholder_tokens(token, num_vec_per_token=num_vec_per_token)
            self.text_encoder.resize_token_embeddings(len(self.tokenizer))
            token_ids = self.tokenizer.encode(token, add_special_tokens=False)

            if embedding_dims > 1:
                for i, token_id in enumerate(token_ids):
                    self.text_encoder.get_input_embeddings().weight.data[token_id] = embedding[i]
            else:
                self.text_encoder.get_input_embeddings().weight.data[token_ids] = embedding

    def _extract_embedding_from_dict(self, embedding_dict: Dict[str, str]) -> torch.Tensor:
        r"""
        Extracts the embedding from the embedding dictionary.
        Arguments:
            embedding_dict (`Dict[str, str]`):
                The embedding dictionary loaded from the embedding path
        Returns:
            embedding (`torch.Tensor`):
                The embedding to be added to the text encoder's embedding matrix
            is_multi_vec_token (`bool`):
                Whether the embedding is a multi-vector token or not
        """
        is_multi_vec_token = False
        # auto1111 embedding case
        if "string_to_param" in embedding_dict:
            embedding_dict = embedding_dict["string_to_param"]
            embedding = embedding_dict["*"]
        else:
            embedding = list(embedding_dict.values())[0]

        if len(embedding.shape) > 1:
            # If the embedding has more than one dimension,
            # We need to ensure the tokenizer is a MultiTokenTokenizer
            # because there is branching logic that depends on that class
            if not isinstance(self.tokenizer, MultiTokenCLIPTokenizer):
                raise ValueError(
                    f"{self.__class__.__name__} requires `self.tokenizer` of type `MultiTokenCLIPTokenizer` for loading embeddings with more than one dimension."
                )
            is_multi_vec_token = True

        return embedding, is_multi_vec_token

    def _extract_token_from_dict(self, embedding_dict: Dict[str, str]) -> str:
        r"""
        Extracts the token from the embedding dictionary.
        Arguments:
            embedding_dict (`Dict[str, str]`):
                The embedding dictionary loaded from the embedding path
        Returns:
            token (`str`):
                The token to be added to the tokenizers' vocabulary
        """
        # auto1111 embedding case
        if "string_to_param" in embedding_dict:
            token = embedding_dict["name"]
            return token

        return list(embedding_dict.keys())[0]

    def _validate_method_call(self, method: Callable):
        r"""
        Validates that the method is being called from a class instance that has the required attributes.
        Arguments:
            method (`function`):
                The class's method being called
        Raises:
            ValueError:
                If the method is being called from a class instance that does not have
                the required attributes, the method will not be callable.
        Returns:
            None
        """
        if not hasattr(self, "tokenizer") or not isinstance(self.tokenizer, PreTrainedTokenizer):
            raise ValueError(
                f"{self.__class__.__name__} requires `self.tokenizer` of type `PreTrainedTokenizer` for calling `{method.__name__}`"
            )

        if not hasattr(self, "text_encoder") or not isinstance(self.text_encoder, PreTrainedModel):
            raise ValueError(
                f"{self.__class__.__name__} requires `self.text_encoder` of type `PreTrainedModel` for calling `{method.__name__}`"
            )

    def _validate_token_update(self, token, allow_replacement=False, is_multi_vec_token=False):
        r"""Validates that the token is not already in the tokenizer's vocabulary.
        Arguments:
            token (`str`):
                The token to be added to the tokenizers' vocabulary
            allow_replacement (`bool`):
                Whether to allow replacement of the token if it already exists in the tokenizer's vocabulary
            is_multi_vec_token (`bool`):
                Whether the embedding is a multi-vector token or not
        Raises:
            ValueError:
                If the token is already in the tokenizer's vocabulary and `allow_replacement` is False.
        Returns:
            None
        """
        if (not is_multi_vec_token and token in self.tokenizer.get_vocab()) or (
            is_multi_vec_token and token in self.tokenizer.token_map
        ):
            if allow_replacement:
                print(
                    f"Token {token} already in tokenizer vocabulary. Overwriting existing token and embedding with the new one."
                )
            else:
                raise ValueError(
                    f"Token {token} already in tokenizer vocabulary. Please choose a different token name."
                )

    
    def expand_textual_inversion_token_ids_if_necessary(self, token_ids: List[int]) -> List[int]:
        pass

class MultiTokenCLIPTokenizer(CLIPTokenizer):
    """Tokenizer for CLIP models that have multi-vector tokens."""

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.token_map = {}

    def add_placeholder_tokens(self, placeholder_token, *args, num_vec_per_token=1, **kwargs):
        r"""Adds placeholder tokens to the tokenizer's vocabulary.
        Arguments:
            placeholder_token (`str`):
                The placeholder token to be added to the tokenizers' vocabulary and token map.
            num_vec_per_token (`int`):
                The number of vectors per token. Defaults to 1.
            *args:
                The arguments to be passed to the tokenizer's `add_tokens` method.
            **kwargs:
                The keyword arguments to be passed to the tokenizer's `add_tokens` method.
        Returns:
            None
        """
        output = []
        if num_vec_per_token == 1:
            self.add_tokens(placeholder_token, *args, **kwargs)
            output.append(placeholder_token)
        else:
            output = []
            for i in range(num_vec_per_token):
                ith_token = placeholder_token + f"_{i}"
                self.add_tokens(ith_token, *args, **kwargs)
                output.append(ith_token)
        # handle cases where there is a new placeholder token that contains the current placeholder token but is larger
        for token in self.token_map:
            if token in placeholder_token:
                raise ValueError(
                    f"The tokenizer already has placeholder token {token} that can get confused with"
                    f" {placeholder_token}keep placeholder tokens independent"
                )
        self.token_map[placeholder_token] = output

    def replace_placeholder_tokens_in_text(self, text, vector_shuffle=False, prop_tokens_to_load=1.0):
        r"""Replaces placeholder tokens in text with the tokens in the token map.
        Opttionally, implements:
            a) vector shuffling (https://github.com/rinongal/textual_inversion/pull/119)where
            shuffling tokens were found to force the model to learn the concepts more descriptively.
            b) proportional token loading so that not every token in the token map is loaded on each call;
            used as part of progressive token loading during training which can improve generalization
            during inference.
        Arguments:
            text (`str`):
                The text to be processed.
            vector_shuffle (`bool`):
                Whether to shuffle the vectors in the token map. Defaults to False.
            prop_tokens_to_load (`float`):
                The proportion of tokens to load from the token map. Defaults to 1.0.
        Returns:
            `str`: The processed text.
        """
        if isinstance(text, list):
            output = []
            for i in range(len(text)):
                output.append(self.replace_placeholder_tokens_in_text(text[i], vector_shuffle=vector_shuffle))
            return output
        for placeholder_token in self.token_map:
            if placeholder_token in text:
                tokens = self.token_map[placeholder_token]
                tokens = tokens[: 1 + int(len(tokens) * prop_tokens_to_load)]
                if vector_shuffle:
                    tokens = copy.copy(tokens)
                    random.shuffle(tokens)
                text = text.replace(placeholder_token, " ".join(tokens))
        return text

    def __call__(self, text, *args, vector_shuffle=False, prop_tokens_to_load=1.0, **kwargs):
        """Wrapper around [`~transformers.tokenization_utils.PreTrainedTokenizerBase.__call__`] method
        but first replace placeholder tokens in text with the tokens in the token map.
        Returns:
            [`~transformers.tokenization_utils_base.BatchEncoding`]
        """
        return super().__call__(
            self.replace_placeholder_tokens_in_text(
                text,
                vector_shuffle=vector_shuffle,
                prop_tokens_to_load=prop_tokens_to_load,
            ),
            *args,
            **kwargs,
        )

    def encode(self, text, *args, vector_shuffle=False, prop_tokens_to_load=1.0, **kwargs):
        """Wrapper around the tokenizer's [`transformers.tokenization_utils.PreTrainedTokenizerBase.encode`] method
        but first replaces placeholder tokens in text with the tokens in the token map.
        Arguments:
            text (`str`):
                The text to be encoded.
            *args:
                The arguments to be passed to the tokenizer's `encode` method.
            vector_shuffle (`bool`):
                Whether to shuffle the vectors in the token map. Defaults to False.
            prop_tokens_to_load (`float`):
                The proportion of tokens to load from the token map. Defaults to 1.0.
            **kwargs:
                The keyword arguments to be passed to the tokenizer's `encode` method.
        Returns:
            List[`int`]: sequence of ids (integer)
        """
        return super().encode(
            self.replace_placeholder_tokens_in_text(
                text,
                vector_shuffle=vector_shuffle,
                prop_tokens_to_load=prop_tokens_to_load,
            ),
            *args,
            **kwargs,
        )