4 files changed, 146 insertions, 13 deletions

diff --git a/dreambooth.py b/dreambooth.py
index 7b61c45..48fc7f2 100644
--- a/dreambooth.py
+++ b/dreambooth.py
@@ -15,14 +15,14 @@ from accelerate import Accelerator
 from accelerate.logging import get_logger
 from accelerate.utils import LoggerType, set_seed
 from diffusers import AutoencoderKL, DDPMScheduler, PNDMScheduler, UNet2DConditionModel
-from schedulers.scheduling_euler_a import EulerAScheduler
 from diffusers.optimization import get_scheduler
 from PIL import Image
 from tqdm.auto import tqdm
 from transformers import CLIPTextModel, CLIPTokenizer
 from slugify import slugify
-from pipelines.stable_diffusion.vlpn_stable_diffusion import VlpnStableDiffusion
 
+from schedulers.scheduling_euler_a import EulerAScheduler
+from pipelines.stable_diffusion.vlpn_stable_diffusion import VlpnStableDiffusion
 from data.csv import CSVDataModule
 
 logger = get_logger(__name__)
@@ -334,7 +334,6 @@ class Checkpointer:
                 beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", skip_prk_steps=True
             ),
         )
-        pipeline.enable_attention_slicing()
         pipeline.save_pretrained(self.output_dir.joinpath("model"))
 
         del unwrapped
@@ -359,7 +358,6 @@ class Checkpointer:
             tokenizer=self.tokenizer,
             scheduler=scheduler,
         ).to(self.accelerator.device)
-        pipeline.enable_attention_slicing()
         pipeline.set_progress_bar_config(dynamic_ncols=True)
 
         train_data = self.datamodule.train_dataloader()
@@ -561,7 +559,6 @@ def main():
                 tokenizer=tokenizer,
                 scheduler=scheduler,
             ).to(accelerator.device)
-            pipeline.enable_attention_slicing()
             pipeline.set_progress_bar_config(dynamic_ncols=True)
 
             with torch.inference_mode():
diff --git a/infer.py b/infer.py
index a542534..70851fd 100644
--- a/infer.py
+++ b/infer.py
@@ -11,8 +11,9 @@ from PIL import Image
 from diffusers import AutoencoderKL, UNet2DConditionModel, PNDMScheduler, DDIMScheduler, LMSDiscreteScheduler
 from transformers import CLIPTextModel, CLIPTokenizer
 from slugify import slugify
-from pipelines.stable_diffusion.vlpn_stable_diffusion import VlpnStableDiffusion
+
 from schedulers.scheduling_euler_a import EulerAScheduler
+from pipelines.stable_diffusion.vlpn_stable_diffusion import VlpnStableDiffusion
 
 
 torch.backends.cuda.matmul.allow_tf32 = True
@@ -235,7 +236,6 @@ def create_pipeline(model, scheduler, embeddings_dir, dtype):
         tokenizer=tokenizer,
         scheduler=scheduler,
     )
-    # pipeline.enable_attention_slicing()
     pipeline.to("cuda")
 
     print("Pipeline loaded.")
diff --git a/models/hypernetwork.py b/models/hypernetwork.py
new file mode 100644
index 0000000..fe8a312
--- /dev/null
+++ b/models/hypernetwork.py
@@ -0,0 +1,138 @@
+import math
+from typing import Dict, Optional, Iterable, List, Tuple, Any
+import copy
+import torch
+import numpy as np
+from torch import nn
+from functorch import make_functional, make_functional_with_buffers
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.modeling_utils import ModelMixin
+
+
+def get_weight_chunk_dims(num_target_parameters: int, num_embeddings: int):
+    weight_chunk_dim = math.ceil(num_target_parameters / num_embeddings)
+    if weight_chunk_dim != 0:
+        remainder = num_target_parameters % weight_chunk_dim
+        if remainder > 0:
+            diff = math.ceil(remainder / weight_chunk_dim)
+            num_embeddings += diff
+    return weight_chunk_dim
+
+
+def count_params(target: ModelMixin):
+    return sum([np.prod(p.size()) for p in target.parameters()])
+
+
+class FunctionalParamVectorWrapper(ModelMixin):
+    """
+    This wraps a module so that it takes params in the forward pass
+    """
+
+    def __init__(self, module: ModelMixin):
+        super().__init__()
+
+        self.custom_buffers = None
+        param_dict = dict(module.named_parameters())
+        self.target_weight_shapes = {k: param_dict[k].size() for k in param_dict}
+
+        try:
+            _functional, self.named_params = make_functional(module)
+        except Exception:
+            _functional, self.named_params, buffers = make_functional_with_buffers(
+                module
+            )
+            self.custom_buffers = buffers
+        self.functional = [_functional]  # remove params from being counted
+
+    def forward(self, param_vector: torch.Tensor, *args, **kwargs):
+        params = []
+        start = 0
+        for p in self.named_params:
+            end = start + np.prod(p.size())
+            params.append(param_vector[start:end].view(p.size()))
+            start = end
+        if self.custom_buffers is not None:
+            return self.functional[0](params, self.custom_buffers, *args, **kwargs)
+        return self.functional[0](params, *args, **kwargs)
+
+
+class Hypernetwork(ModelMixin, ConfigMixin):
+    @register_to_config
+    def __init__(
+        self,
+        target_network: ModelMixin,
+        num_target_parameters: Optional[int] = None,
+        embedding_dim: int = 100,
+        num_embeddings: int = 3,
+        weight_chunk_dim: Optional[int] = None,
+    ):
+        super().__init__()
+
+        self._target = FunctionalParamVectorWrapper(target_network)
+
+        self.target_weight_shapes = self._target.target_weight_shapes
+
+        self.num_target_parameters = num_target_parameters
+
+        self.embedding_dim = embedding_dim
+        self.num_embeddings = num_embeddings
+        self.weight_chunk_dim = weight_chunk_dim
+
+        self.embedding_module = self.make_embedding_module()
+        self.weight_generator = self.make_weight_generator()
+
+    def make_embedding_module(self) -> nn.Module:
+        return nn.Embedding(self.num_embeddings, self.embedding_dim)
+
+    def make_weight_generator(self) -> nn.Module:
+        return nn.Linear(self.embedding_dim, self.weight_chunk_dim)
+
+    def generate_params(
+        self, inp: Iterable[Any] = []
+    ) -> Tuple[torch.Tensor, Dict[str, Any]]:
+        embedding = self.embedding_module(
+            torch.arange(self.num_embeddings, device=self.device)
+        )
+        generated_params = self.weight_generator(embedding).view(-1)
+        return generated_params, {"embedding": embedding}
+
+    def forward(
+        self,
+        inp: Iterable[Any] = [],
+        *args,
+        **kwargs,
+    ):
+        generated_params, aux_output = self.generate_params(inp, *args, **kwargs)
+
+        assert generated_params.shape[-1] >= self.num_target_parameters
+
+        return self._target(generated_params, *inp)
+
+    @property
+    def device(self) -> torch.device:
+        return self._target.device
+
+    @classmethod
+    def from_target(
+        cls,
+        target_network: ModelMixin,
+        num_target_parameters: Optional[int] = None,
+        embedding_dim: int = 8,
+        num_embeddings: int = 3,
+        weight_chunk_dim: Optional[int] = None,
+        *args,
+        **kwargs,
+    ):
+        if num_target_parameters is None:
+            num_target_parameters = count_params(target_network)
+        if weight_chunk_dim is None:
+            weight_chunk_dim = get_weight_chunk_dims(num_target_parameters, num_embeddings)
+        return cls(
+            target_network=target_network,
+            num_target_parameters=num_target_parameters,
+            embedding_dim=embedding_dim,
+            num_embeddings=num_embeddings,
+            weight_chunk_dim=weight_chunk_dim,
+            *args,
+            **kwargs,
+        )
diff --git a/textual_inversion.py b/textual_inversion.py
index 09871d4..e641cab 100644
--- a/textual_inversion.py
+++ b/textual_inversion.py
@@ -16,14 +16,14 @@ from accelerate import Accelerator
 from accelerate.logging import get_logger
 from accelerate.utils import LoggerType, set_seed
 from diffusers import AutoencoderKL, DDPMScheduler, UNet2DConditionModel
-from schedulers.scheduling_euler_a import EulerAScheduler
 from diffusers.optimization import get_scheduler
 from PIL import Image
 from tqdm.auto import tqdm
 from transformers import CLIPTextModel, CLIPTokenizer
 from slugify import slugify
-from pipelines.stable_diffusion.vlpn_stable_diffusion import VlpnStableDiffusion
 
+from schedulers.scheduling_euler_a import EulerAScheduler
+from pipelines.stable_diffusion.vlpn_stable_diffusion import VlpnStableDiffusion
 from data.csv import CSVDataModule
 
 logger = get_logger(__name__)
@@ -388,7 +388,6 @@ class Checkpointer:
             tokenizer=self.tokenizer,
             scheduler=scheduler,
         ).to(self.accelerator.device)
-        pipeline.enable_attention_slicing()
         pipeline.set_progress_bar_config(dynamic_ncols=True)
 
         train_data = self.datamodule.train_dataloader()
@@ -518,8 +517,8 @@ def main():
     if args.gradient_checkpointing:
         unet.enable_gradient_checkpointing()
 
-    slice_size = unet.config.attention_head_dim // 2
-    unet.set_attention_slice(slice_size)
+    # slice_size = unet.config.attention_head_dim // 2
+    # unet.set_attention_slice(slice_size)
 
     # Resize the token embeddings as we are adding new special tokens to the tokenizer
     text_encoder.resize_token_embeddings(len(tokenizer))
@@ -639,7 +638,6 @@ def main():
                 tokenizer=tokenizer,
                 scheduler=scheduler,
             ).to(accelerator.device)
-            pipeline.enable_attention_slicing()
             pipeline.set_progress_bar_config(dynamic_ncols=True)
 
             with torch.inference_mode():