| 392 | |
| 393 | |
| 394 | def encode_token_weights(self, token_weight_pairs): |
| 395 | to_encode = list() |
| 396 | max_token_len = 0 |
| 397 | has_weights = False |
| 398 | for x in token_weight_pairs: |
| 399 | tokens = list(map(lambda a: a[0], x)) |
| 400 | max_token_len = max(len(tokens), max_token_len) |
| 401 | has_weights = has_weights or not all(map(lambda a: a[1] == 1.0, x)) |
| 402 | to_encode.append(tokens) |
| 403 | |
| 404 | sections = len(to_encode) |
| 405 | if has_weights or sections == 0: |
| 406 | if hasattr(self, "gen_empty_tokens"): |
| 407 | to_encode.append(self.gen_empty_tokens(self.special_tokens, max_token_len)) |
| 408 | else: |
| 409 | to_encode.append(comfy.sd1_clip.gen_empty_tokens(self.special_tokens, max_token_len)) |
| 410 | |
| 411 | o = self.encode(to_encode) |
| 412 | out, pooled = o[:2] |
| 413 | |
| 414 | # if pooled is not None: |
| 415 | # first_pooled = pooled[0:1].to(model_management.intermediate_device()) |
| 416 | # else: |
| 417 | # first_pooled = pooled |
| 418 | assert pooled is None |
| 419 | first_pooled = None |
| 420 | |
| 421 | output = [] |
| 422 | for k in range(0, sections): |
| 423 | z = out[k:k+1] |
| 424 | if has_weights: |
| 425 | z_empty = out[-1] |
| 426 | for i in range(len(z)): |
| 427 | for j in range(len(z[i])): |
| 428 | weight = token_weight_pairs[k][j][1] |
| 429 | if weight != 1.0: |
| 430 | z[i][j] = (z[i][j] - z_empty[j]) * weight + z_empty[j] |
| 431 | output.append(z) |
| 432 | |
| 433 | if (len(output) == 0): |
| 434 | r = (out[-1:].to(model_management.intermediate_device()), first_pooled) |
| 435 | else: |
| 436 | r = (torch.cat(output, dim=0).to(model_management.intermediate_device()), first_pooled) |
| 437 | |
| 438 | if len(o) > 2: |
| 439 | extra = {} |
| 440 | for k in o[2]: |
| 441 | v = o[2][k] |
| 442 | extra[k] = v |
| 443 | |
| 444 | r = r + (extra,) |
| 445 | return r |
| 446 | |
| 447 | # Handle batch of different prompts. |
| 448 | comfy.sd1_clip.ClipTokenWeightEncoder.encode_token_weights = encode_token_weights |