Optimized embedding normalization

1 files changed, 2 insertions, 5 deletions

diff --git a/models/clip/embeddings.py b/models/clip/embeddings.py
index 1cc59d9..6c41c33 100644
--- a/models/clip/embeddings.py
+++ b/models/clip/embeddings.py
@@ -3,7 +3,6 @@ from pathlib import Path
 
 import torch
 import torch.nn as nn
-import torch.nn.functional as F
 
 from safetensors import safe_open
 from safetensors.torch import save_file
@@ -104,10 +103,8 @@ class ManagedCLIPTextEmbeddings(CLIPTextEmbeddings):
             return
 
         w = self.temp_token_embedding.weight
-        pre_norm = w[self.temp_token_ids, :].norm(dim=-1, keepdim=True)
-        w[self.temp_token_ids] = F.normalize(
-            w[self.temp_token_ids, :], dim=-1
-        ) * (pre_norm + lambda_ * (target - pre_norm))
+        norm = w[self.temp_token_ids, :].norm(dim=-1, keepdim=True)
+        w[self.temp_token_ids].add_((w[self.temp_token_ids] / norm.clamp_min(1e-12)) * lambda_ * (target - norm))
 
     def forward(
         self,