add a function for extracting the necessary parameters for fine-tuning, directly off the stable diffusion (or any text-to-<modality>) instance being operated on

Author: lucidrains

README.md

@@ -48,7 +48,8 @@ wrapped_to_values = Rank1EditModule(
 
 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_indices = torch.randint(0, 77, (4,))
+concept_indices = torch.randint(0, 77, (4,))        # index where the concept or superclass concept token is in the sequence
+key_pad_mask = torch.ones(4, 77).bool()
 
 keys = wrapped_to_keys(
     text_enc,
@@ -81,7 +82,7 @@ from torch import nn
 
 from perfusion_pytorch import EmbeddingWrapper
 
-embed = nn.Embedding(49407, 512) # open clip embedding, somewhere in the module tree of stable diffusion
+embed = nn.Embedding(49408, 512) # open clip embedding, somewhere in the module tree of stable diffusion
 
 # wrap it, and will automatically create a new concept for learning, based on the superclass embed string
 

perfusion_pytorch/__init__.py

@@ -13,5 +13,6 @@
 
 from perfusion_pytorch.save_load import (
     save,
-    load
+    load,
+    get_finetune_parameters
 )

perfusion_pytorch/embedding.py

@@ -71,7 +71,7 @@ def __init__(
         self.num_concepts = num_concepts
         self.concepts = nn.Parameter(torch.zeros(num_concepts, dim))
 
-        assert exists(superclass_embed_id) ^ exists(superclass_string), 'either superclass embed id is given, or the superclass string'
+        assert not (exists(superclass_embed_id) and exists(superclass_string)), 'either superclass embed id is given, or the superclass string'
 
         self.pad_id = tokenizer_pad_id
         self.tokenize = None
@@ -130,8 +130,8 @@ def forward(
             if not isinstance(concept_id, tuple):
                 concept_id = (concept_id,)
 
-            assert len(concept_id) == 1, 'can only train or inference on single concepts if passing in list of superclass strings'
-            assert self.num_concepts == 1
+            assert not self.training or len(concept_id) == 1, 'can only train or inference on single concepts if passing in list of superclass strings'
+            assert not self.training or self.num_concepts == 1
 
         if is_bearable(x, List[str]):
             inferred_concept_id = self.concept_embed_ids[0]
@@ -221,6 +221,8 @@ def merge_embedding_wrappers(
         num_concepts = total_concepts
     )
 
+    merged_concepts.eval()
+
     concepts = torch.cat(tuple(embed.concepts.data for embed in embeds), dim = 0)
 
     merged_concepts.concepts = nn.Parameter(concepts)

perfusion_pytorch/save_load.py

@@ -9,9 +9,24 @@
 from perfusion_pytorch.embedding import EmbeddingWrapper
 from perfusion_pytorch.perfusion import Rank1EditModule
 
+# helper functions
+
 def exists(val):
     return val is not None
 
+# function that automatically finds all the parameters necessary for fine tuning
+
+@beartype
+def get_finetune_parameters(text_image_model: Module):
+    params = []
+    for module in text_image_model.modules():
+        if isinstance(module, (EmbeddingWrapper, Rank1EditModule)):
+            params.extend(module.parameters())
+
+    return params
+
+# saving and loading the necessary extra finetuned params
+
 @beartype
 def save(
     text_image_model: Module,

setup.py

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