take into account some feedback in email with author

Author: lucidrains

README.md

@@ -23,10 +23,10 @@ import torch
 from perfusion_pytorch import Rank1EditModule
 from torch import nn
 
-to_keys = nn.Linear(512, 1024, bias = False)
-to_values = nn.Linear(512, 1024, bias = False)
+to_keys = nn.Linear(768, 320, bias = False)
+to_values = nn.Linear(768, 320, bias = False)
 
-input_covariance = torch.randn(512, 512)
+input_covariance = torch.randn(768, 768)
 
 wrapped_to_keys = Rank1EditModule(
     to_keys,
@@ -41,10 +41,9 @@ wrapped_to_values = Rank1EditModule(
     num_finetune_prompts = 32
 )
 
-text_enc = torch.randn(4, 256, 512)                  # regular input
-text_enc_with_superclass = torch.randn(4, 256, 512)  # init_input in algorithm 1, for key-locking
-concept_ids = torch.randint(0, 256, (4,))
-
+text_enc = torch.randn(4, 77, 768)                  # regular input
+text_enc_with_superclass = torch.randn(4, 77, 768)  # init_input in algorithm 1, for key-locking
+concept_ids = torch.randint(0, 77, (4,))
 
 keys = wrapped_to_keys(text_enc, text_enc_with_superclass, concept_ids)
 values = wrapped_to_values(text_enc, text_enc_with_superclass, concept_ids)

perfusion_pytorch/perfusion.py

@@ -26,7 +26,7 @@ def __init__(
         *,
         num_finetune_prompts: int,
         C: Tensor,                  # covariance of input, precomputed from 100K laion text
-        text_seq_len: int = 256,
+        text_seq_len: int = 77,
         is_key_proj: bool = False,
         input_decay = 0.99,
         train_beta = 0.75,
@@ -92,16 +92,16 @@ def forward(
         concept_text_enc = text_enc[batch_indices, concept_indices]
         concept_text_enc = rearrange(concept_text_enc, 'b 1 d -> b d')
 
-        superclass_text_enc = text_enc_with_superclass[batch_indices, concept_indices]
-        superclass_text_enc = rearrange(superclass_text_enc, 'b 1 d -> b d')
-
         # take care of initializing with superclass prompt
         # for key-locking - this assumes stable diffusion was modified so text encoder takes in a prompt with both the <concept> as well as <superclass> - it seems this also has the limitation that <superclass> must be one token
 
-        text_enc_with_superclass_output = einsum('b n i, o i -> b n o', text_enc_with_superclass, weights)
+        superclass_text_enc = text_enc_with_superclass[batch_indices, concept_indices]
+        superclass_text_enc = rearrange(superclass_text_enc, 'b 1 d -> b d')
+
+        superclass_text_enc_output = einsum('b i, o i -> b o', superclass_text_enc, weights)
 
         if self.is_key_proj:
-            text_enc_with_superclass_output = text_enc_with_superclass_output.detach()
+            superclass_text_enc_output = superclass_text_enc_output.detach()
 
         # only during training do they exponentially smooth
 
@@ -112,7 +112,7 @@ def forward(
 
         # make it easier to match with paper
 
-        i, o, W = online_estimated_concept_enc, text_enc_with_superclass_output, weights
+        i, o, W = online_estimated_concept_enc, superclass_text_enc_output, weights
 
         # main contribution eq (3)
 
@@ -131,4 +131,4 @@ def forward(
 
         W_em_orthogonal_term = text_enc_output - (sim * concept_output / i_energy)
 
-        return W_em_orthogonal_term + sigmoid_term * o
+        return W_em_orthogonal_term + sigmoid_term * rearrange(o, 'b d -> b 1 d')

setup.py

@@ -3,7 +3,7 @@
 setup(
   name = 'perfusion-pytorch',
   packages = find_packages(exclude=[]),
-  version = '0.0.6',
+  version = '0.0.7',
   license='MIT',
   description = 'Perfusion - Pytorch',
   author = 'Phil Wang',