From ad3ae441081155dcd4fde805279e5082ca264695 Mon Sep 17 00:00:00 2001
From: Fampai <unknown>
Date: Sun, 9 Oct 2022 04:32:40 -0400
Subject: [PATCH] Updated code for legibility

---
 modules/sd_hijack.py | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/modules/sd_hijack.py b/modules/sd_hijack.py
index 4a2d21531..7793d25be 100644
--- a/modules/sd_hijack.py
+++ b/modules/sd_hijack.py
@@ -284,8 +284,11 @@ class FrozenCLIPEmbedderWithCustomWords(torch.nn.Module):
 
         tmp = -opts.CLIP_stop_at_last_layers
         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)
+        if tmp < -1:
+            z = outputs.hidden_states[tmp]
+            z = self.wrapped.transformer.text_model.final_layer_norm(z)
+        else:
+            z = outputs.last_hidden_state
 
         # 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]