Signed-off-by: Vladimir Mandic <mandic00@live.com>
This commit is contained in:
Vladimir Mandic
2024-10-01 20:02:21 -04:00
parent f9b5a83b49
commit 9a46d381cc
8 changed files with 47 additions and 40 deletions
+27 -25
View File
@@ -34,10 +34,12 @@ def compel_hijack(self, token_ids: torch.Tensor,
attention_mask: typing.Optional[torch.Tensor] = None) -> torch.Tensor:
needs_hidden_states = self.returned_embeddings_type != 1
try: # can crash in ATen/native/cuda/Indexing since position_ids are corrupt so index lookup fails, but its not compel specific, happens with fixed attention as well
sd_models.move_model(self.text_encoder, devices.device)
text_encoder_output = self.text_encoder(token_ids, attention_mask, output_hidden_states=needs_hidden_states, return_dict=True)
except Exception as e: # its a non-recoverable error as cuda state is corrupt
shared.log.error(f'TE: class={self.text_encoder.__class__} device={self.text_encoder.device} dtype={self.text_encoder.dtype} {e}')
errors.display(e, 'TE:')
return None
if not needs_hidden_states:
return text_encoder_output.last_hidden_state
@@ -174,7 +176,7 @@ def encode_prompts(pipe, p, prompts: list, negative_prompts: list, steps: int, c
):
shared.log.warning(f"Prompt parser not supported: {pipe.__class__.__name__}")
return
elif shared.opts.sd_textencoder_cache and cache.get('model_type', None) == shared.sd_model_type and params_match and False:
elif shared.opts.sd_textencoder_cache and cache.get('model_type', None) == shared.sd_model_type and params_match:
p.prompt_embeds = cache.get('prompt_embeds', None)
p.positive_pooleds = cache.get('positive_pooleds', None)
p.negative_embeds = cache.get('negative_embeds', None)
@@ -190,10 +192,14 @@ def encode_prompts(pipe, p, prompts: list, negative_prompts: list, steps: int, c
pipe.maybe_free_model_hooks()
devices.torch_gc()
fix_position_ids(pipe)
prompt_embeds, positive_pooleds, negative_embeds, negative_pooleds = [], [], [], []
p.prompt_embeds = []
p.positive_pooleds = []
p.negative_embeds = []
p.negative_pooleds = []
p.scheduled_prompt = False
last_prompt, last_negative = None, None
for prompt, negative in zip(prompts, negative_prompts):
prompt_embeds, positive_pooleds, negative_embeds, negative_pooleds = [], [], [], []
prompt_embed, positive_pooled, negative_embed, negative_pooled = None, None, None, None
if last_prompt == prompt and last_negative == negative:
prompt_embeds.append(prompt_embeds[-1])
@@ -205,11 +211,7 @@ def encode_prompts(pipe, p, prompts: list, negative_prompts: list, steps: int, c
continue
positive_schedule, scheduled = get_prompt_schedule(prompt, steps)
negative_schedule, neg_scheduled = get_prompt_schedule(negative, steps)
p.scheduled_prompt = scheduled or neg_scheduled
p.prompt_embeds = []
p.positive_pooleds = []
p.negative_embeds = []
p.negative_pooleds = []
p.scheduled_prompt = p.scheduled_prompt or scheduled or neg_scheduled
for i in range(max(len(positive_schedule), len(negative_schedule))):
positive_prompt = positive_schedule[i % len(positive_schedule)]
@@ -228,25 +230,25 @@ def encode_prompts(pipe, p, prompts: list, negative_prompts: list, steps: int, c
negative_pooleds.append(negative_pooled)
last_prompt, last_negative = prompt, negative
def fix_length(embeds):
max_len = max([e.shape[1] for e in embeds if e is not None])
for i, e in enumerate(embeds):
if e is not None and e.shape[1] < max_len:
expanded = torch.zeros((e.shape[0], max_len, e.shape[2]), device=e.device, dtype=e.dtype)
expanded[:, :e.shape[1], :] = e
embeds[i] = expanded
return torch.cat(embeds, dim=0).to(devices.device, dtype=devices.dtype)
def fix_length(embeds):
max_len = max([e.shape[1] for e in embeds if e is not None])
for i, e in enumerate(embeds):
if e is not None and e.shape[1] < max_len:
expanded = torch.zeros((e.shape[0], max_len, e.shape[2]), device=e.device, dtype=e.dtype)
expanded[:, :e.shape[1], :] = e
embeds[i] = expanded
return torch.cat(embeds, dim=0).to(devices.device, dtype=devices.dtype)
if len(prompt_embeds) > 0:
p.prompt_embeds.append(fix_length(prompt_embeds))
if len(negative_embeds) > 0:
p.negative_embeds.append(fix_length(negative_embeds))
if len(positive_pooleds) > 0:
p.positive_pooleds.append(fix_length(positive_pooleds))
if len(negative_pooleds) > 0:
p.negative_pooleds.append(fix_length(negative_pooleds))
if len(prompt_embeds) > 0:
p.prompt_embeds.append(fix_length(prompt_embeds))
if len(negative_embeds) > 0:
p.negative_embeds.append(fix_length(negative_embeds))
if len(positive_pooleds) > 0:
p.positive_pooleds.append(fix_length(positive_pooleds))
if len(negative_pooleds) > 0:
p.negative_pooleds.append(fix_length(negative_pooleds))
if shared.opts.sd_textencoder_cache and p.batch_size == 1:
if p.batch_size == 1:
cache.update({
'prompt_embeds': p.prompt_embeds,
'negative_embeds': p.negative_embeds,