3 files changed, 125 insertions, 215 deletions

diff --git a/training/common.py b/training/functional.py
index 5d1e3f9..2d81eca 100644
--- a/training/common.py
+++ b/training/functional.py
@@ -16,19 +16,14 @@ from pipelines.stable_diffusion.vlpn_stable_diffusion import VlpnStableDiffusion
 from models.clip.embeddings import ManagedCLIPTextEmbeddings, patch_managed_embeddings
 from models.clip.util import get_extended_embeddings
 from models.clip.tokenizer import MultiCLIPTokenizer
-from training.util import AverageMeter, CheckpointerBase
+from training.util import AverageMeter
+from trainer.base import Checkpointer
 
 
-def noop(*args, **kwards):
-    pass
-
-
-def noop_ctx(*args, **kwards):
-    return nullcontext()
-
-
-def noop_on_log():
-    return {}
+def const(result=None):
+    def fn(*args, **kwargs):
+        return result
+    return fn
 
 
 def generate_class_images(
@@ -210,7 +205,7 @@ def train_loop(
     optimizer: torch.optim.Optimizer,
     lr_scheduler: torch.optim.lr_scheduler._LRScheduler,
     model: torch.nn.Module,
-    checkpointer: CheckpointerBase,
+    checkpointer: Checkpointer,
     train_dataloader: DataLoader,
     val_dataloader: DataLoader,
     loss_step: Union[Callable[[int, Any], Tuple[Any, Any, int]], Callable[[int, Any, bool], Tuple[Any, Any, int]]],
@@ -218,11 +213,11 @@ def train_loop(
     checkpoint_frequency: int = 50,
     global_step_offset: int = 0,
     num_epochs: int = 100,
-    on_log: Callable[[], dict[str, Any]] = noop_on_log,
-    on_train: Callable[[int], _GeneratorContextManager] = noop_ctx,
-    on_before_optimize: Callable[[int], None] = noop,
-    on_after_optimize: Callable[[float], None] = noop,
-    on_eval: Callable[[], _GeneratorContextManager] = noop_ctx
+    on_log: Callable[[], dict[str, Any]] = const({}),
+    on_train: Callable[[int], _GeneratorContextManager] = const(nullcontext()),
+    on_before_optimize: Callable[[int], None] = const(),
+    on_after_optimize: Callable[[float], None] = const(),
+    on_eval: Callable[[], _GeneratorContextManager] = const(nullcontext())
 ):
     num_training_steps_per_epoch = math.ceil(len(train_dataloader) / accelerator.gradient_accumulation_steps)
     num_val_steps_per_epoch = len(val_dataloader)
diff --git a/training/lora.py b/training/lora.py
deleted file mode 100644
index 3857d78..0000000
--- a/training/lora.py
+++ /dev/null
@@ -1,107 +0,0 @@
-import torch
-import torch.nn as nn
-
-from diffusers import ModelMixin, ConfigMixin
-from diffusers.configuration_utils import register_to_config
-from diffusers.models.cross_attention import CrossAttention
-from diffusers.utils.import_utils import is_xformers_available
-
-
-if is_xformers_available():
-    import xformers
-    import xformers.ops
-else:
-    xformers = None
-
-
-class LoRALinearLayer(nn.Module):
-    def __init__(self, in_features, out_features, rank=4):
-        super().__init__()
-
-        if rank > min(in_features, out_features):
-            raise ValueError(
-                f"LoRA rank {rank} must be less or equal than {min(in_features, out_features)}"
-            )
-
-        self.lora_down = nn.Linear(in_features, rank, bias=False)
-        self.lora_up = nn.Linear(rank, out_features, bias=False)
-        self.scale = 1.0
-
-        nn.init.normal_(self.lora_down.weight, std=1 / rank)
-        nn.init.zeros_(self.lora_up.weight)
-
-    def forward(self, hidden_states):
-        down_hidden_states = self.lora_down(hidden_states)
-        up_hidden_states = self.lora_up(down_hidden_states)
-
-        return up_hidden_states
-
-
-class LoRACrossAttnProcessor(nn.Module):
-    def __init__(self, hidden_size, cross_attention_dim=None, rank=4):
-        super().__init__()
-
-        self.to_q_lora = LoRALinearLayer(hidden_size, hidden_size)
-        self.to_k_lora = LoRALinearLayer(cross_attention_dim or hidden_size, hidden_size)
-        self.to_v_lora = LoRALinearLayer(cross_attention_dim or hidden_size, hidden_size)
-        self.to_out_lora = LoRALinearLayer(hidden_size, hidden_size)
-
-    def __call__(self, attn: CrossAttention, hidden_states, encoder_hidden_states=None, attention_mask=None, scale=1.0):
-        batch_size, sequence_length, _ = hidden_states.shape
-        attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length)
-
-        query = attn.to_q(hidden_states) + scale * self.to_q_lora(hidden_states)
-        query = attn.head_to_batch_dim(query)
-
-        encoder_hidden_states = encoder_hidden_states if encoder_hidden_states is not None else hidden_states
-
-        key = attn.to_k(encoder_hidden_states) + scale * self.to_k_lora(encoder_hidden_states)
-        value = attn.to_v(encoder_hidden_states) + scale * self.to_v_lora(encoder_hidden_states)
-
-        key = attn.head_to_batch_dim(key)
-        value = attn.head_to_batch_dim(value)
-
-        attention_probs = attn.get_attention_scores(query, key, attention_mask)
-        hidden_states = torch.bmm(attention_probs, value)
-        hidden_states = attn.batch_to_head_dim(hidden_states)
-
-        # linear proj
-        hidden_states = attn.to_out[0](hidden_states) + scale * self.to_out_lora(hidden_states)
-        # dropout
-        hidden_states = attn.to_out[1](hidden_states)
-
-        return hidden_states
-
-
-class LoRAXFormersCrossAttnProcessor(nn.Module):
-    def __init__(self, hidden_size, cross_attention_dim, rank=4):
-        super().__init__()
-
-        self.to_q_lora = LoRALinearLayer(hidden_size, hidden_size)
-        self.to_k_lora = LoRALinearLayer(cross_attention_dim or hidden_size, hidden_size)
-        self.to_v_lora = LoRALinearLayer(cross_attention_dim or hidden_size, hidden_size)
-        self.to_out_lora = LoRALinearLayer(hidden_size, hidden_size)
-
-    def __call__(self, attn: CrossAttention, hidden_states, encoder_hidden_states=None, attention_mask=None, scale=1.0):
-        batch_size, sequence_length, _ = hidden_states.shape
-        attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length)
-
-        query = attn.to_q(hidden_states) + scale * self.to_q_lora(hidden_states)
-        query = attn.head_to_batch_dim(query).contiguous()
-
-        encoder_hidden_states = encoder_hidden_states if encoder_hidden_states is not None else hidden_states
-
-        key = attn.to_k(encoder_hidden_states) + scale * self.to_k_lora(encoder_hidden_states)
-        value = attn.to_v(encoder_hidden_states) + scale * self.to_v_lora(encoder_hidden_states)
-
-        key = attn.head_to_batch_dim(key).contiguous()
-        value = attn.head_to_batch_dim(value).contiguous()
-
-        hidden_states = xformers.ops.memory_efficient_attention(query, key, value, attn_bias=attention_mask)
-
-        # linear proj
-        hidden_states = attn.to_out[0](hidden_states) + scale * self.to_out_lora(hidden_states)
-        # dropout
-        hidden_states = attn.to_out[1](hidden_states)
-
-        return hidden_states
diff --git a/training/util.py b/training/util.py
index 781cf04..a292edd 100644
--- a/training/util.py
+++ b/training/util.py
@@ -1,12 +1,40 @@
 from pathlib import Path
 import json
 import copy
-import itertools
-from typing import Iterable, Optional
+from typing import Iterable, Union
 from contextlib import contextmanager
 
 import torch
-from PIL import Image
+
+from transformers import CLIPTextModel
+from diffusers import AutoencoderKL, DDPMScheduler, UNet2DConditionModel, DPMSolverMultistepScheduler
+
+from pipelines.stable_diffusion.vlpn_stable_diffusion import VlpnStableDiffusion
+from models.clip.tokenizer import MultiCLIPTokenizer
+from models.clip.embeddings import ManagedCLIPTextEmbeddings, patch_managed_embeddings
+
+
+class TrainingStrategy():
+    @property
+    def main_model(self) -> torch.nn.Module:
+        ...
+
+    @contextmanager
+    def on_train(self, epoch: int):
+        yield
+
+    @contextmanager
+    def on_eval(self):
+        yield
+
+    def on_before_optimize(self, epoch: int):
+        ...
+
+    def on_after_optimize(self, lr: float):
+        ...
+
+    def on_log():
+        return {}
 
 
 def save_args(basepath: Path, args, extra={}):
@@ -16,113 +44,107 @@ def save_args(basepath: Path, args, extra={}):
         json.dump(info, f, indent=4)
 
 
-def make_grid(images, rows, cols):
-    w, h = images[0].size
-    grid = Image.new('RGB', size=(cols*w, rows*h))
-    for i, image in enumerate(images):
-        grid.paste(image, box=(i % cols*w, i//cols*h))
-    return grid
+def generate_class_images(
+    accelerator,
+    text_encoder,
+    vae,
+    unet,
+    tokenizer,
+    scheduler,
+    data_train,
+    sample_batch_size,
+    sample_image_size,
+    sample_steps
+):
+    missing_data = [item for item in data_train if not item.class_image_path.exists()]
 
+    if len(missing_data) == 0:
+        return
 
-class AverageMeter:
-    def __init__(self, name=None):
-        self.name = name
-        self.reset()
+    batched_data = [
+        missing_data[i:i+sample_batch_size]
+        for i in range(0, len(missing_data), sample_batch_size)
+    ]
 
-    def reset(self):
-        self.sum = self.count = self.avg = 0
+    pipeline = VlpnStableDiffusion(
+        text_encoder=text_encoder,
+        vae=vae,
+        unet=unet,
+        tokenizer=tokenizer,
+        scheduler=scheduler,
+    ).to(accelerator.device)
+    pipeline.set_progress_bar_config(dynamic_ncols=True)
 
-    def update(self, val, n=1):
-        self.sum += val * n
-        self.count += n
-        self.avg = self.sum / self.count
+    with torch.inference_mode():
+        for batch in batched_data:
+            image_name = [item.class_image_path for item in batch]
+            prompt = [item.cprompt for item in batch]
+            nprompt = [item.nprompt for item in batch]
 
+            images = pipeline(
+                prompt=prompt,
+                negative_prompt=nprompt,
+                height=sample_image_size,
+                width=sample_image_size,
+                num_inference_steps=sample_steps
+            ).images
 
-class CheckpointerBase:
-    def __init__(
-        self,
-        train_dataloader,
-        val_dataloader,
-        output_dir: Path,
-        sample_steps: int = 20,
-        sample_guidance_scale: float = 7.5,
-        sample_image_size: int = 768,
-        sample_batches: int = 1,
-        sample_batch_size: int = 1,
-        seed: Optional[int] = None
-    ):
-        self.train_dataloader = train_dataloader
-        self.val_dataloader = val_dataloader
-        self.output_dir = output_dir
-        self.sample_image_size = sample_image_size
-        self.sample_steps = sample_steps
-        self.sample_guidance_scale = sample_guidance_scale
-        self.sample_batches = sample_batches
-        self.sample_batch_size = sample_batch_size
-        self.seed = seed if seed is not None else torch.random.seed()
+            for i, image in enumerate(images):
+                image.save(image_name[i])
 
-    @torch.no_grad()
-    def checkpoint(self, step: int, postfix: str):
-        pass
+    del pipeline
 
-    @torch.inference_mode()
-    def save_samples(self, pipeline, step: int):
-        samples_path = Path(self.output_dir).joinpath("samples")
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
 
-        generator = torch.Generator(device=pipeline.device).manual_seed(self.seed)
 
-        grid_cols = min(self.sample_batch_size, 4)
-        grid_rows = (self.sample_batches * self.sample_batch_size) // grid_cols
+def get_models(pretrained_model_name_or_path: str):
+    tokenizer = MultiCLIPTokenizer.from_pretrained(pretrained_model_name_or_path, subfolder='tokenizer')
+    text_encoder = CLIPTextModel.from_pretrained(pretrained_model_name_or_path, subfolder='text_encoder')
+    vae = AutoencoderKL.from_pretrained(pretrained_model_name_or_path, subfolder='vae')
+    unet = UNet2DConditionModel.from_pretrained(pretrained_model_name_or_path, subfolder='unet')
+    noise_scheduler = DDPMScheduler.from_pretrained(pretrained_model_name_or_path, subfolder='scheduler')
+    sample_scheduler = DPMSolverMultistepScheduler.from_pretrained(
+        pretrained_model_name_or_path, subfolder='scheduler')
 
-        for pool, data, gen in [
-            ("stable", self.val_dataloader, generator),
-            ("val", self.val_dataloader, None),
-            ("train", self.train_dataloader, None)
-        ]:
-            all_samples = []
-            file_path = samples_path.joinpath(pool, f"step_{step}.jpg")
-            file_path.parent.mkdir(parents=True, exist_ok=True)
+    embeddings = patch_managed_embeddings(text_encoder)
 
-            batches = list(itertools.islice(itertools.cycle(data), self.sample_batch_size * self.sample_batches))
-            prompt_ids = [
-                prompt
-                for batch in batches
-                for prompt in batch["prompt_ids"]
-            ]
-            nprompt_ids = [
-                prompt
-                for batch in batches
-                for prompt in batch["nprompt_ids"]
-            ]
+    return tokenizer, text_encoder, vae, unet, noise_scheduler, sample_scheduler, embeddings
 
-            for i in range(self.sample_batches):
-                start = i * self.sample_batch_size
-                end = (i + 1) * self.sample_batch_size
-                prompt = prompt_ids[start:end]
-                nprompt = nprompt_ids[start:end]
 
-                samples = pipeline(
-                    prompt=prompt,
-                    negative_prompt=nprompt,
-                    height=self.sample_image_size,
-                    width=self.sample_image_size,
-                    generator=gen,
-                    guidance_scale=self.sample_guidance_scale,
-                    num_inference_steps=self.sample_steps,
-                    output_type='pil'
-                ).images
+def add_placeholder_tokens(
+    tokenizer: MultiCLIPTokenizer,
+    embeddings: ManagedCLIPTextEmbeddings,
+    placeholder_tokens: list[str],
+    initializer_tokens: list[str],
+    num_vectors: Union[list[int], int]
+):
+    initializer_token_ids = [
+        tokenizer.encode(token, add_special_tokens=False)
+        for token in initializer_tokens
+    ]
+    placeholder_token_ids = tokenizer.add_multi_tokens(placeholder_tokens, num_vectors)
 
-                all_samples += samples
+    embeddings.resize(len(tokenizer))
 
-                del samples
+    for (placeholder_token_id, initializer_token_id) in zip(placeholder_token_ids, initializer_token_ids):
+        embeddings.add_embed(placeholder_token_id, initializer_token_id)
 
-            image_grid = make_grid(all_samples, grid_rows, grid_cols)
-            image_grid.save(file_path, quality=85)
+    return placeholder_token_ids, initializer_token_ids
 
-            del all_samples
-            del image_grid
 
-        del generator
+class AverageMeter:
+    def __init__(self, name=None):
+        self.name = name
+        self.reset()
+
+    def reset(self):
+        self.sum = self.count = self.avg = 0
+
+    def update(self, val, n=1):
+        self.sum += val * n
+        self.count += n
+        self.avg = self.sum / self.count
 
 
 # Adapted from torch-ema https://github.com/fadel/pytorch_ema/blob/master/torch_ema/ema.py#L14