almost there

Author: lucidrains

README.md

@@ -75,10 +75,11 @@ values = wrapped_to_values(text_enc)
 
 ## Todo
 
-- [ ] offer a magic function that automatically tries to wire up the cross attention by looking for appropriately named `nn.Linear` and auto-inferring which ones are keys or values
 - [ ] show example in readme for inference with multiple concepts
 - [ ] review multiple concepts
+- [ ] automatically infer where keys and values projection are if not specified for the `make_key_value_proj_rank1_edit_modules_` function
 
+- [x] offer a function that wires up the cross attention
 - [x] handle multiple concepts in one prompt at inference - summation of the sigmoid term + outputs
     - [x] accept multiple concept indices
 - [x] offer a way to combine separately learned concepts from multiple `Rank1EditModule` into one for inference

perfusion_pytorch/__init__.py

@@ -2,7 +2,8 @@
     Rank1EditModule,
     calculate_input_covariance,
     loss_fn_weighted_by_mask,
-    merge_rank1_edit_modules
+    merge_rank1_edit_modules,
+    make_key_value_proj_rank1_edit_modules_
 )
 
 from perfusion_pytorch.embedding import (

perfusion_pytorch/perfusion.py

@@ -388,3 +388,26 @@ def merge_rank1_edit_modules(
     merged_module.register_buffer('initted', torch.ones(total_concepts, 1).bool())
 
     return merged_module
+
+# function for wiring up the cross attention
+
+@beartype
+def make_key_value_proj_rank1_edit_modules_(
+    cross_attention: nn.Module,
+    *,
+    input_covariance: Tensor,
+    key_proj_name: str,
+    value_proj_name: str,
+    **rank1_edit_module_kwargs
+):
+    linear_key = getattr(cross_attention, key_proj_name, None)
+    linear_values = getattr(cross_attention, value_proj_name, None)
+
+    assert isinstance(linear_key, nn.Linear), f'{key_proj_name} must point to where the keys projection is (ex. self.to_keys = nn.Linear(in, out, bias = False) -> key_proj_name = "to_keys")'
+    assert isinstance(linear_values, nn.Linear), f'{value_proj_name} must point to where the values projection is (ex. self.to_keys = nn.Linear(in, out, bias = False) -> value_proj_name = "to_values")'
+
+    rank1_edit_module_keys = Rank1EditModule(linear_key, input_covariance = input_covariance, is_key_proj = True, **rank1_edit_module_kwargs)
+    rank1_edit_module_values = Rank1EditModule(linear_values, input_covariance = input_covariance, is_key_proj = False, **rank1_edit_module_kwargs)
+
+    setattr(cross_attention, key_proj_name, rank1_edit_module_keys)
+    setattr(cross_attention, value_proj_name, rank1_edit_module_values)

setup.py

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