| 175 | |
| 176 | |
| 177 | class SDTokenizer: |
| 178 | def __init__(self, max_length=77, pad_with_end=True, tokenizer=None, has_start_token=True, pad_to_max_length=True, min_length=None): |
| 179 | self.tokenizer = tokenizer |
| 180 | self.max_length = max_length |
| 181 | self.min_length = min_length |
| 182 | empty = self.tokenizer('')["input_ids"] |
| 183 | if has_start_token: |
| 184 | self.tokens_start = 1 |
| 185 | self.start_token = empty[0] |
| 186 | self.end_token = empty[1] |
| 187 | else: |
| 188 | self.tokens_start = 0 |
| 189 | self.start_token = None |
| 190 | self.end_token = empty[0] |
| 191 | self.pad_with_end = pad_with_end |
| 192 | self.pad_to_max_length = pad_to_max_length |
| 193 | vocab = self.tokenizer.get_vocab() |
| 194 | self.inv_vocab = {v: k for k, v in vocab.items()} |
| 195 | self.max_word_length = 8 |
| 196 | |
| 197 | |
| 198 | def tokenize_with_weights(self, text:str): |
| 199 | """Tokenize the text, with weight values - presume 1.0 for all and ignore other features here. The details aren't relevant for a reference impl, and weights themselves has weak effect on SD3.""" |
| 200 | if self.pad_with_end: |
| 201 | pad_token = self.end_token |
| 202 | else: |
| 203 | pad_token = 0 |
| 204 | batch = [] |
| 205 | if self.start_token is not None: |
| 206 | batch.append((self.start_token, 1.0)) |
| 207 | to_tokenize = text.replace("\n", " ").split(' ') |
| 208 | to_tokenize = [x for x in to_tokenize if x != ""] |
| 209 | for word in to_tokenize: |
| 210 | batch.extend([(t, 1) for t in self.tokenizer(word)["input_ids"][self.tokens_start:-1]]) |
| 211 | batch.append((self.end_token, 1.0)) |
| 212 | if self.pad_to_max_length: |
| 213 | batch.extend([(pad_token, 1.0)] * (self.max_length - len(batch))) |
| 214 | if self.min_length is not None and len(batch) < self.min_length: |
| 215 | batch.extend([(pad_token, 1.0)] * (self.min_length - len(batch))) |
| 216 | return [batch] |
| 217 | |
| 218 | |
| 219 | class SDXLClipGTokenizer(SDTokenizer): |