4 files changed, 132 insertions, 64 deletions

diff --git a/models/clip/embeddings.py b/models/clip/embeddings.py
index c9c788c..1e21965 100644
--- a/models/clip/embeddings.py
+++ b/models/clip/embeddings.py
@@ -31,41 +31,15 @@ def resize_embedding(old_embedding: nn.Embedding, new_num_embeddings: int, initi
     return new_embedding
 
 
-class OverlayLinear(nn.Module):
-    def __init__(self, in_features, out_features, rank=4):
-        super().__init__()
-
-        if rank > min(in_features, out_features):
-            raise ValueError(f"Rank {rank} must be less or equal than {min(in_features, out_features)}")
-
-        self.rank = rank
-        self.down = nn.Linear(in_features, rank, bias=False)
-        self.up = nn.Linear(rank, out_features, bias=False)
-        self.reset()
-
-    def reset(self):
-        nn.init.normal_(self.down.weight, std=1 / self.rank)
-        nn.init.zeros_(self.up.weight)
-
-    def forward(self, hidden_states):
-        orig_dtype = hidden_states.dtype
-        dtype = self.down.weight.dtype
-
-        down_hidden_states = self.down(hidden_states.to(dtype))
-        up_hidden_states = self.up(down_hidden_states)
-
-        return up_hidden_states.to(orig_dtype)
-
-
 class ManagedCLIPTextEmbeddings(CLIPTextEmbeddings):
-    def __init__(self, config: CLIPTextConfig, embeddings: CLIPTextEmbeddings, rank: int = 128):
+    def __init__(self, config: CLIPTextConfig, embeddings: CLIPTextEmbeddings, alpha: float = 1.0, rank: int = 4):
         super().__init__(config)
 
         self.token_embedding = embeddings.token_embedding
         self.position_embedding = embeddings.position_embedding
         self.initializer_factor = config.initializer_factor
+        self.alpha = alpha
 
-        self.overlay = OverlayLinear(self.token_embedding.embedding_dim, self.token_embedding.embedding_dim, rank)
         self.temp_token_embedding = nn.Embedding(
             self.token_embedding.num_embeddings,
             self.token_embedding.embedding_dim,
@@ -75,9 +49,6 @@ class ManagedCLIPTextEmbeddings(CLIPTextEmbeddings):
         self.temp_token_embedding.weight.data = self.token_embedding.weight.data.clone().detach()
         self.temp_token_ids = torch.tensor([], dtype=torch.long)
 
-    def reset_overlay(self):
-        self.overlay.reset()
-
     def resize(self, size: int):
         self.temp_token_embedding = resize_embedding(self.temp_token_embedding, size, self.initializer_factor)
         self.token_embedding = resize_embedding(self.token_embedding, size, self.initializer_factor)
@@ -125,9 +96,7 @@ class ManagedCLIPTextEmbeddings(CLIPTextEmbeddings):
         save_file({"embed": self.get_embed(input_ids)}, filename)
 
     def persist(self):
-        embeds = self.temp_token_embedding.weight.data[self.temp_token_ids]
-        self.token_embedding.weight.data[self.temp_token_ids] = embeds + self.overlay(embeds)
-        self.overlay.reset()
+        self.token_embedding.weight.data[self.temp_token_ids] = self.temp_token_embedding.weight.data[self.temp_token_ids]
         self.temp_token_ids = torch.tensor([], dtype=torch.long)
 
     def get_embed(self, input_ids: Union[list[int], torch.LongTensor]):
@@ -135,11 +104,8 @@ class ManagedCLIPTextEmbeddings(CLIPTextEmbeddings):
             input_ids = torch.tensor(input_ids, device=self.token_embedding.weight.device, dtype=torch.long)
 
         embeds = self.token_embedding(input_ids)
-
         mask = torch.isin(input_ids, self.temp_token_ids.to(input_ids.device))
-
-        temp_embeds = self.temp_token_embedding(input_ids[mask])
-        embeds[mask] = temp_embeds + self.overlay(temp_embeds)
+        embeds[mask] = self.temp_token_embedding(input_ids[mask])
 
         return embeds
 
diff --git a/train_ti.py b/train_ti.py
index 26ac384..5482326 100644
--- a/train_ti.py
+++ b/train_ti.py
@@ -1,7 +1,6 @@
 import argparse
 import datetime
 import logging
-import itertools
 from functools import partial
 from pathlib import Path
 import math
@@ -308,6 +307,26 @@ def parse_args():
         help="Minimum learning rate in the lr scheduler."
     )
     parser.add_argument(
+        "--use_ema",
+        action="store_true",
+        help="Whether to use EMA model."
+    )
+    parser.add_argument(
+        "--ema_inv_gamma",
+        type=float,
+        default=1.0
+    )
+    parser.add_argument(
+        "--ema_power",
+        type=float,
+        default=4/5
+    )
+    parser.add_argument(
+        "--ema_max_decay",
+        type=float,
+        default=0.9999
+    )
+    parser.add_argument(
         "--optimizer",
         type=str,
         default="dadan",
@@ -334,7 +353,7 @@ def parse_args():
     parser.add_argument(
         "--adam_weight_decay",
         type=float,
-        default=1e-2,
+        default=0,
         help="Weight decay to use."
     )
     parser.add_argument(
@@ -432,6 +451,23 @@ def parse_args():
         help="The weight of prior preservation loss."
     )
     parser.add_argument(
+        "--use_emb_decay",
+        action="store_true",
+        help="Whether to use embedding decay."
+    )
+    parser.add_argument(
+        "--emb_decay_target",
+        default=0.4,
+        type=float,
+        help="Embedding decay target."
+    )
+    parser.add_argument(
+        "--emb_decay",
+        default=1e2,
+        type=float,
+        help="Embedding decay factor."
+    )
+    parser.add_argument(
         "--noise_timesteps",
         type=int,
         default=1000,
@@ -696,6 +732,13 @@ def main():
         sample_scheduler=sample_scheduler,
         checkpoint_output_dir=checkpoint_output_dir,
         gradient_checkpointing=args.gradient_checkpointing,
+        use_emb_decay=args.use_emb_decay,
+        emb_decay_target=args.emb_decay_target,
+        emb_decay=args.emb_decay,
+        use_ema=args.use_ema,
+        ema_inv_gamma=args.ema_inv_gamma,
+        ema_power=args.ema_power,
+        ema_max_decay=args.ema_max_decay,
         sample_batch_size=args.sample_batch_size,
         sample_num_batches=args.sample_batches,
         sample_num_steps=args.sample_steps,
@@ -757,10 +800,7 @@ def main():
             sample_frequency = math.ceil(num_train_epochs * (sample_frequency / args.num_train_steps))
 
         optimizer = create_optimizer(
-            itertools.chain(
-                text_encoder.text_model.embeddings.temp_token_embedding.parameters(),
-                text_encoder.text_model.embeddings.overlay.parameters(),
-            ),
+            text_encoder.text_model.embeddings.temp_token_embedding.parameters(),
             lr=args.learning_rate,
         )
 
diff --git a/training/functional.py b/training/functional.py
index 7104a88..bd8cbad 100644
--- a/training/functional.py
+++ b/training/functional.py
@@ -524,7 +524,7 @@ def train_loop(
 
             lr = lr_scheduler.get_last_lr()[0]
             if torch.is_tensor(lr):
-                lr = lr.item
+                lr = lr.item()
 
             lrs.append(lr)
 
diff --git a/training/strategy/ti.py b/training/strategy/ti.py
index 1b5adab..677f5a3 100644
--- a/training/strategy/ti.py
+++ b/training/strategy/ti.py
@@ -1,6 +1,6 @@
 from typing import Optional
 from functools import partial
-from contextlib import contextmanager
+from contextlib import contextmanager, nullcontext
 from pathlib import Path
 
 import torch
@@ -13,6 +13,7 @@ from diffusers import AutoencoderKL, UNet2DConditionModel, DPMSolverMultistepSch
 from slugify import slugify
 
 from models.clip.tokenizer import MultiCLIPTokenizer
+from training.util import EMAModel
 from training.functional import TrainingStrategy, TrainingCallbacks, save_samples
 
 
@@ -31,6 +32,13 @@ def textual_inversion_strategy_callbacks(
     placeholder_tokens: list[str],
     placeholder_token_ids: list[list[int]],
     gradient_checkpointing: bool = False,
+    use_emb_decay: bool = False,
+    emb_decay_target: float = 0.4,
+    emb_decay: float = 1e-2,
+    use_ema: bool = False,
+    ema_inv_gamma: float = 1.0,
+    ema_power: int = 1,
+    ema_max_decay: float = 0.9999,
     sample_batch_size: int = 1,
     sample_num_batches: int = 1,
     sample_num_steps: int = 20,
@@ -63,8 +71,27 @@ def textual_inversion_strategy_callbacks(
         image_size=sample_image_size,
     )
 
+    if use_ema:
+        ema_embeddings = EMAModel(
+            text_encoder.text_model.embeddings.temp_token_embedding.parameters(),
+            inv_gamma=ema_inv_gamma,
+            power=ema_power,
+            max_value=ema_max_decay,
+        )
+        ema_embeddings.to(accelerator.device)
+    else:
+        ema_embeddings = None
+
+    def ema_context():
+        if ema_embeddings is not None:
+            return ema_embeddings.apply_temporary(
+                text_encoder.text_model.embeddings.temp_token_embedding.parameters()
+            )
+        else:
+            return nullcontext()
+
     def on_accum_model():
-        return text_encoder.text_model.embeddings
+        return text_encoder.text_model.embeddings.temp_token_embedding
 
     @contextmanager
     def on_train(epoch: int):
@@ -74,36 +101,68 @@ def textual_inversion_strategy_callbacks(
     @contextmanager
     def on_eval():
         tokenizer.eval()
-        yield
+
+        with ema_context():
+            yield
+
+    @torch.no_grad()
+    def on_before_optimize(lr: float, epoch: int):
+        if use_emb_decay:
+            w = text_encoder.text_model.embeddings.temp_token_embedding.weight
+            return torch.all(w.grad == 0, dim=1)
+
+    @torch.no_grad()
+    def on_after_optimize(zero_ids, lr: float):
+        if ema_embeddings is not None:
+            ema_embeddings.step(text_encoder.text_model.embeddings.temp_token_embedding.parameters())
+
+        if use_emb_decay:
+            lambda_ = emb_decay * lr
+
+            if lambda_ != 0:
+                w = text_encoder.text_model.embeddings.temp_token_embedding.weight
+
+                mask = torch.ones(w.shape[0], dtype=torch.bool)
+                mask[zero_ids] = False
+
+                norm = w[mask, :].norm(dim=-1, keepdim=True)
+                w[mask].add_((w[mask] / norm.clamp_min(1e-12)) * lambda_ * (emb_decay_target - norm))
+
+    def on_log():
+        if ema_embeddings is not None:
+            return {"ema_decay": ema_embeddings.decay}
+        return {}
 
     @torch.no_grad()
     def on_checkpoint(step, postfix):
         print(f"Saving checkpoint for step {step}...")
 
-        for (token, ids) in zip(placeholder_tokens, placeholder_token_ids):
-            text_encoder.text_model.embeddings.save_embed(
-                ids,
-                checkpoint_output_dir / f"{slugify(token)}_{step}_{postfix}.bin"
-            )
+        with ema_context():
+            for (token, ids) in zip(placeholder_tokens, placeholder_token_ids):
+                text_encoder.text_model.embeddings.save_embed(
+                    ids,
+                    checkpoint_output_dir / f"{slugify(token)}_{step}_{postfix}.bin"
+                )
 
     @torch.no_grad()
     def on_sample(step):
-        unet_ = accelerator.unwrap_model(unet, keep_fp32_wrapper=True)
-        text_encoder_ = accelerator.unwrap_model(text_encoder, keep_fp32_wrapper=True)
+        with ema_context():
+            unet_ = accelerator.unwrap_model(unet, keep_fp32_wrapper=True)
+            text_encoder_ = accelerator.unwrap_model(text_encoder, keep_fp32_wrapper=True)
 
-        orig_unet_dtype = unet_.dtype
-        orig_text_encoder_dtype = text_encoder_.dtype
+            orig_unet_dtype = unet_.dtype
+            orig_text_encoder_dtype = text_encoder_.dtype
 
-        unet_.to(dtype=weight_dtype)
-        text_encoder_.to(dtype=weight_dtype)
+            unet_.to(dtype=weight_dtype)
+            text_encoder_.to(dtype=weight_dtype)
 
-        save_samples_(step=step, unet=unet_, text_encoder=text_encoder_)
+            save_samples_(step=step, unet=unet_, text_encoder=text_encoder_)
 
-        unet_.to(dtype=orig_unet_dtype)
-        text_encoder_.to(dtype=orig_text_encoder_dtype)
+            unet_.to(dtype=orig_unet_dtype)
+            text_encoder_.to(dtype=orig_text_encoder_dtype)
 
-        del unet_
-        del text_encoder_
+            del unet_
+            del text_encoder_
 
         if torch.cuda.is_available():
             torch.cuda.empty_cache()
@@ -112,6 +171,9 @@ def textual_inversion_strategy_callbacks(
         on_accum_model=on_accum_model,
         on_train=on_train,
         on_eval=on_eval,
+        on_before_optimize=on_before_optimize,
+        on_after_optimize=on_after_optimize,
+        on_log=on_log,
         on_checkpoint=on_checkpoint,
         on_sample=on_sample,
     )