mirror of
https://github.com/lllyasviel/stable-diffusion-webui-forge.git
synced 2026-07-21 21:01:24 +08:00
153 lines
4.8 KiB
Python
153 lines
4.8 KiB
Python
# https://github.com/comfyanonymous/ComfyUI/blob/v0.3.64/comfy/sd1_clip.py
|
|
# https://github.com/comfyanonymous/ComfyUI/blob/v0.3.64/comfy/text_encoders/wan.py
|
|
|
|
import torch
|
|
|
|
from backend import memory_management
|
|
from backend.args import dynamic_args
|
|
from backend.text_processing import emphasis, parsing
|
|
from modules.shared import opts
|
|
|
|
|
|
class PromptChunk:
|
|
def __init__(self):
|
|
self.tokens = []
|
|
self.multipliers = []
|
|
|
|
|
|
class UMT5TextProcessingEngine:
|
|
def __init__(self, text_encoder, tokenizer, min_length=512):
|
|
self.emphasis = emphasis.get_current_option(opts.emphasis)()
|
|
|
|
self.text_encoder = text_encoder.transformer
|
|
self.tokenizer = tokenizer
|
|
self.device = memory_management.text_encoder_device()
|
|
|
|
self.max_length = 99999999
|
|
self.min_length = min_length
|
|
|
|
empty = self.tokenizer("")["input_ids"]
|
|
self.tokens_start = 0
|
|
self.tokens_end = -1
|
|
self.end_token = empty[0]
|
|
self.pad_token = 0
|
|
|
|
def tokenize(self, texts):
|
|
return self.tokenizer(texts)["input_ids"]
|
|
|
|
def process_attn_mask(self, tokens):
|
|
attention_masks = []
|
|
|
|
for x in tokens:
|
|
attention_mask = []
|
|
eos = False
|
|
|
|
for y in x:
|
|
if isinstance(y, int):
|
|
attention_mask.append(0 if eos else 1)
|
|
if not eos and int(y) == self.end_token:
|
|
eos = True
|
|
|
|
attention_masks.append(attention_mask)
|
|
|
|
return torch.tensor(attention_masks, dtype=torch.long, device=self.device)
|
|
|
|
def encode_with_transformers(self, tokens, attention_mask):
|
|
tokens = tokens.to(self.device)
|
|
return self.text_encoder(input_ids=tokens, attention_mask=attention_mask)
|
|
|
|
def tokenize_line(self, line):
|
|
parsed = parsing.parse_prompt_attention(line, self.emphasis.name)
|
|
|
|
tokenized = self.tokenize([text[self.tokens_start : self.tokens_end] for text, _ in parsed])
|
|
|
|
chunks = []
|
|
chunk = PromptChunk()
|
|
token_count = 0
|
|
|
|
def next_chunk():
|
|
nonlocal token_count
|
|
nonlocal chunk
|
|
|
|
chunk.tokens.append(self.end_token)
|
|
chunk.multipliers.append(1.0)
|
|
|
|
current_chunk_length = len(chunk.tokens)
|
|
token_count += current_chunk_length
|
|
|
|
if current_chunk_length < self.min_length:
|
|
chunk.tokens.extend([self.pad_token] * (self.min_length - current_chunk_length))
|
|
chunk.multipliers.extend([1.0] * (self.min_length - current_chunk_length))
|
|
|
|
chunks.append(chunk)
|
|
chunk = PromptChunk()
|
|
|
|
for tokens, (text, weight) in zip(tokenized, parsed):
|
|
if text == "BREAK" and weight == -1:
|
|
next_chunk()
|
|
continue
|
|
|
|
position = 0
|
|
while position < len(tokens):
|
|
token = tokens[position]
|
|
chunk.tokens.append(token)
|
|
chunk.multipliers.append(weight)
|
|
position += 1
|
|
|
|
if chunk.tokens or not chunks:
|
|
next_chunk()
|
|
|
|
return chunks, token_count
|
|
|
|
def __call__(self, texts):
|
|
self.emphasis = emphasis.get_current_option(opts.emphasis)()
|
|
if any(emphasis.uses_emphasis(x) for x in texts):
|
|
dynamic_args.last_extra_generation_params["Emphasis"] = self.emphasis.name
|
|
|
|
zs = []
|
|
cache = {}
|
|
|
|
for line in texts:
|
|
if line in cache:
|
|
line_z_values = cache[line]
|
|
else:
|
|
chunks, _ = self.tokenize_line(line)
|
|
line_z_values = []
|
|
|
|
# pad all chunks to length of longest chunk
|
|
max_tokens = 0
|
|
for chunk in chunks:
|
|
max_tokens = max(len(chunk.tokens), max_tokens)
|
|
|
|
for chunk in chunks:
|
|
tokens = chunk.tokens
|
|
multipliers = chunk.multipliers
|
|
|
|
remaining_count = max_tokens - len(tokens)
|
|
if remaining_count > 0:
|
|
tokens += [self.id_pad] * remaining_count
|
|
multipliers += [1.0] * remaining_count
|
|
|
|
z = self.process_tokens([tokens], [multipliers])[0]
|
|
line_z_values.append(z)
|
|
cache[line] = line_z_values
|
|
|
|
zs.extend(line_z_values)
|
|
|
|
return torch.stack(zs)
|
|
|
|
def process_tokens(self, batch_tokens, batch_multipliers):
|
|
tokens = torch.asarray(batch_tokens)
|
|
|
|
attention_mask = self.process_attn_mask(batch_tokens)
|
|
z = self.encode_with_transformers(tokens, attention_mask)
|
|
z *= attention_mask.unsqueeze(-1).float()
|
|
|
|
self.emphasis.tokens = batch_tokens
|
|
self.emphasis.multipliers = torch.asarray(batch_multipliers).to(z)
|
|
self.emphasis.z = z
|
|
self.emphasis.after_transformers()
|
|
z = self.emphasis.z
|
|
|
|
return z
|