Added ability to ignore last n layers in FrozenCLIPEmbedder

This commit is contained in:
Fampai
2022-10-08 14:28:22 -04:00
committed by AUTOMATIC1111
parent b458fa48fe
commit 1371d7608b
2 changed files with 10 additions and 2 deletions

View File

@@ -281,8 +281,15 @@ class FrozenCLIPEmbedderWithCustomWords(torch.nn.Module):
remade_batch_tokens_of_same_length = [x + [self.wrapped.tokenizer.eos_token_id] * (target_token_count - len(x)) for x in remade_batch_tokens]
tokens = torch.asarray(remade_batch_tokens_of_same_length).to(device)
outputs = self.wrapped.transformer(input_ids=tokens, position_ids=position_ids)
z = outputs.last_hidden_state
tmp = -opts.CLIP_ignore_last_layers
if (opts.CLIP_ignore_last_layers == 0):
outputs = self.wrapped.transformer(input_ids=tokens, position_ids=position_ids)
z = outputs.last_hidden_state
else:
outputs = self.wrapped.transformer(input_ids=tokens, position_ids=position_ids, output_hidden_states=tmp)
z = outputs.hidden_states[tmp]
z = self.wrapped.transformer.text_model.final_layer_norm(z)
# restoring original mean is likely not correct, but it seems to work well to prevent artifacts that happen otherwise
batch_multipliers_of_same_length = [x + [1.0] * (target_token_count - len(x)) for x in batch_multipliers]