6 files changed, 541 insertions, 32 deletions

diff --git a/training/common.py b/training/common.py
index 180396e..73ce814 100644
--- a/training/common.py
+++ b/training/common.py
@@ -1,46 +1,77 @@
 import math
+from pathlib import Path
 from contextlib import _GeneratorContextManager, nullcontext
-from typing import Callable, Any, Tuple, Union
+from typing import Callable, Any, Tuple, Union, Literal, Optional, NamedTuple
+import datetime
+import logging
 
 import torch
 import torch.nn.functional as F
 from torch.utils.data import DataLoader
 
 from accelerate import Accelerator
-from transformers import CLIPTokenizer, CLIPTextModel
-from diffusers import AutoencoderKL, DDPMScheduler, UNet2DConditionModel
+from accelerate.utils import LoggerType, set_seed
+from transformers import CLIPTextModel
+from diffusers import AutoencoderKL, DDPMScheduler, UNet2DConditionModel, DPMSolverMultistepScheduler
 from diffusers.optimization import get_scheduler as get_scheduler_, get_cosine_with_hard_restarts_schedule_with_warmup
 
 from tqdm.auto import tqdm
+from slugify import slugify
 
+from data.csv import VlpnDataModule, VlpnDataItem
+from util import load_embeddings_from_dir
 from pipelines.stable_diffusion.vlpn_stable_diffusion import VlpnStableDiffusion
+from models.clip.embeddings import patch_managed_embeddings
 from models.clip.util import get_extended_embeddings
+from models.clip.tokenizer import MultiCLIPTokenizer
 from training.optimization import get_one_cycle_schedule
 from training.util import AverageMeter, CheckpointerBase
 
 
+class TrainingSetup(NamedTuple):
+    accelerator: Accelerator
+    tokenizer: MultiCLIPTokenizer
+    text_encoder: CLIPTextModel
+    vae: AutoencoderKL
+    unet: UNet2DConditionModel
+    noise_scheduler: DDPMScheduler
+    checkpoint_scheduler: DPMSolverMultistepScheduler
+    optimizer_class: Callable
+    learning_rate: float
+    weight_dtype: torch.dtype
+    output_dir: Path
+    seed: int
+    train_dataloader: DataLoader
+    val_dataloader: DataLoader
+    placeholder_token: list[str]
+    placeholder_token_ids: list[list[int]]
+
+
 def noop(*args, **kwards):
     pass
 
 
+def noop_ctx(*args, **kwards):
+    return nullcontext()
+
+
 def noop_on_log():
     return {}
 
 
 def get_scheduler(
     id: str,
-    min_lr: float,
-    lr: float,
-    warmup_func: str,
-    annealing_func: str,
-    warmup_exp: int,
-    annealing_exp: int,
-    cycles: int,
-    train_epochs: int,
-    warmup_epochs: int,
     optimizer: torch.optim.Optimizer,
     num_training_steps_per_epoch: int,
     gradient_accumulation_steps: int,
+    min_lr: float = 0.04,
+    warmup_func: str = "cos",
+    annealing_func: str = "cos",
+    warmup_exp: int = 1,
+    annealing_exp: int = 1,
+    cycles: int = 1,
+    train_epochs: int = 100,
+    warmup_epochs: int = 10,
 ):
     num_training_steps_per_epoch = math.ceil(
         num_training_steps_per_epoch / gradient_accumulation_steps
@@ -49,8 +80,6 @@ def get_scheduler(
     num_warmup_steps = warmup_epochs * num_training_steps_per_epoch
 
     if id == "one_cycle":
-        min_lr = 0.04 if min_lr is None else min_lr / lr
-
         lr_scheduler = get_one_cycle_schedule(
             optimizer=optimizer,
             num_training_steps=num_training_steps,
@@ -133,6 +162,196 @@ def generate_class_images(
             torch.cuda.empty_cache()
 
 
+def train_setup(
+    output_dir: str,
+    project: str,
+    pretrained_model_name_or_path: str,
+    learning_rate: float,
+    data_file: str,
+    gradient_accumulation_steps: int = 1,
+    mixed_precision: Literal["no", "fp16", "bf16"] = "no",
+    seed: Optional[int] = None,
+    vector_shuffle: Union[bool, Literal["all", "trailing", "leading", "between", "off"]] = "auto",
+    vector_dropout: float = 0.1,
+    gradient_checkpointing: bool = True,
+    embeddings_dir: Optional[str] = None,
+    placeholder_token: list[str] = [],
+    initializer_token: list[str] = [],
+    num_vectors: int = 1,
+    scale_lr: bool = False,
+    use_8bit_adam: bool = False,
+    train_batch_size: int = 1,
+    class_image_dir: Optional[str] = None,
+    num_class_images: int = 0,
+    resolution: int = 768,
+    num_buckets: int = 0,
+    progressive_buckets: bool = False,
+    bucket_step_size: int = 64,
+    bucket_max_pixels: Optional[int] = None,
+    tag_dropout: float = 0.1,
+    tag_shuffle: bool = True,
+    data_template: str = "template",
+    valid_set_size: Optional[int] = None,
+    valid_set_repeat: int = 1,
+    data_filter: Optional[Callable[[VlpnDataItem], bool]] = None,
+    sample_batch_size: int = 1,
+    sample_image_size: int = 768,
+    sample_steps: int = 20,
+) -> TrainingSetup:
+    now = datetime.datetime.now().strftime("%Y-%m-%dT%H-%M-%S")
+    output_dir = Path(output_dir).joinpath(slugify(project), now)
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    accelerator = Accelerator(
+        log_with=LoggerType.TENSORBOARD,
+        logging_dir=f"{output_dir}",
+        gradient_accumulation_steps=gradient_accumulation_steps,
+        mixed_precision=mixed_precision
+    )
+
+    logging.basicConfig(filename=output_dir.joinpath("log.txt"), level=logging.DEBUG)
+
+    seed = seed or (torch.random.seed() >> 32)
+    set_seed(seed)
+
+    # Load the tokenizer and add the placeholder token as a additional special token
+    tokenizer = MultiCLIPTokenizer.from_pretrained(pretrained_model_name_or_path, subfolder='tokenizer')
+    tokenizer.set_use_vector_shuffle(vector_shuffle)
+    tokenizer.set_dropout(vector_dropout)
+
+    # Load models and create wrapper for stable diffusion
+    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')
+    checkpoint_scheduler = DPMSolverMultistepScheduler.from_pretrained(
+        pretrained_model_name_or_path, subfolder='scheduler')
+
+    vae.enable_slicing()
+    vae.set_use_memory_efficient_attention_xformers(True)
+    unet.set_use_memory_efficient_attention_xformers(True)
+
+    if gradient_checkpointing:
+        unet.enable_gradient_checkpointing()
+        text_encoder.gradient_checkpointing_enable()
+
+    embeddings = patch_managed_embeddings(text_encoder)
+
+    if embeddings_dir is not None:
+        embeddings_dir = Path(embeddings_dir)
+        if not embeddings_dir.exists() or not embeddings_dir.is_dir():
+            raise ValueError("--embeddings_dir must point to an existing directory")
+
+        added_tokens, added_ids = load_embeddings_from_dir(tokenizer, embeddings, embeddings_dir)
+        print(f"Added {len(added_tokens)} tokens from embeddings dir: {list(zip(added_tokens, added_ids))}")
+
+    # Convert the initializer_token, placeholder_token to ids
+    initializer_token_ids = [
+        tokenizer.encode(token, add_special_tokens=False)
+        for token in initializer_token
+    ]
+
+    placeholder_token_ids = tokenizer.add_multi_tokens(placeholder_token, num_vectors)
+    embeddings.resize(len(tokenizer))
+
+    for (new_id, init_ids) in zip(placeholder_token_ids, initializer_token_ids):
+        embeddings.add_embed(new_id, init_ids)
+
+    init_ratios = [
+        f"{len(init_ids)} / {len(new_id)}"
+        for new_id, init_ids in zip(placeholder_token_ids, initializer_token_ids)
+    ]
+
+    print(f"Added {len(placeholder_token_ids)} new tokens: {list(zip(placeholder_token, placeholder_token_ids, init_ratios))}")
+
+    vae.requires_grad_(False)
+    unet.requires_grad_(False)
+    text_encoder.requires_grad_(False)
+
+    if scale_lr:
+        learning_rate = (
+            learning_rate * gradient_accumulation_steps *
+            train_batch_size * accelerator.num_processes
+        )
+
+    # Use 8-bit Adam for lower memory usage or to fine-tune the model in 16GB GPUs
+    if use_8bit_adam:
+        try:
+            import bitsandbytes as bnb
+        except ImportError:
+            raise ImportError("To use 8-bit Adam, please install the bitsandbytes library: `pip install bitsandbytes`.")
+
+        optimizer_class = bnb.optim.AdamW8bit
+    else:
+        optimizer_class = torch.optim.AdamW
+
+    weight_dtype = torch.float32
+    if mixed_precision == "fp16":
+        weight_dtype = torch.float16
+    elif mixed_precision == "bf16":
+        weight_dtype = torch.bfloat16
+
+    datamodule = VlpnDataModule(
+        data_file=data_file,
+        batch_size=train_batch_size,
+        tokenizer=tokenizer,
+        class_subdir=class_image_dir,
+        num_class_images=num_class_images,
+        size=resolution,
+        num_buckets=num_buckets,
+        progressive_buckets=progressive_buckets,
+        bucket_step_size=bucket_step_size,
+        bucket_max_pixels=bucket_max_pixels,
+        dropout=tag_dropout,
+        shuffle=tag_shuffle,
+        template_key=data_template,
+        valid_set_size=valid_set_size,
+        valid_set_repeat=valid_set_repeat,
+        seed=seed,
+        filter=data_filter,
+        dtype=weight_dtype
+    )
+    datamodule.setup()
+
+    train_dataloader = datamodule.train_dataloader
+    val_dataloader = datamodule.val_dataloader
+
+    train_dataloader, val_dataloader = accelerator.prepare(train_dataloader, val_dataloader)
+
+    if num_class_images != 0:
+        generate_class_images(
+            accelerator,
+            text_encoder,
+            vae,
+            unet,
+            tokenizer,
+            checkpoint_scheduler,
+            datamodule.data_train,
+            sample_batch_size,
+            sample_image_size,
+            sample_steps
+        )
+
+    return TrainingSetup(
+        accelerator=accelerator,
+        tokenizer=tokenizer,
+        text_encoder=text_encoder,
+        vae=vae,
+        unet=unet,
+        noise_scheduler=noise_scheduler,
+        checkpoint_scheduler=checkpoint_scheduler,
+        optimizer_class=optimizer_class,
+        learning_rate=learning_rate,
+        output_dir=output_dir,
+        weight_dtype=weight_dtype,
+        seed=seed,
+        train_dataloader=train_dataloader,
+        val_dataloader=val_dataloader,
+        placeholder_token=placeholder_token,
+        placeholder_token_ids=placeholder_token_ids
+    )
+
+
 def loss_step(
     vae: AutoencoderKL,
     noise_scheduler: DDPMScheduler,
@@ -221,15 +440,14 @@ def train_loop(
     sample_steps: int = 20,
     checkpoint_frequency: int = 50,
     global_step_offset: int = 0,
-    gradient_accumulation_steps: int = 1,
     num_epochs: int = 100,
     on_log: Callable[[], dict[str, Any]] = noop_on_log,
-    on_train: Callable[[], _GeneratorContextManager] = nullcontext,
-    on_before_optimize: Callable[[], None] = noop,
+    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] = nullcontext
+    on_eval: Callable[[], _GeneratorContextManager] = noop_ctx
 ):
-    num_training_steps_per_epoch = math.ceil(len(train_dataloader) / gradient_accumulation_steps)
+    num_training_steps_per_epoch = math.ceil(len(train_dataloader) / accelerator.gradient_accumulation_steps)
     num_val_steps_per_epoch = len(val_dataloader)
 
     num_training_steps = num_training_steps_per_epoch * num_epochs
@@ -273,14 +491,14 @@ def train_loop(
 
             model.train()
 
-            with on_train():
+            with on_train(epoch):
                 for step, batch in enumerate(train_dataloader):
                     with accelerator.accumulate(model):
                         loss, acc, bsz = loss_step(step, batch)
 
                         accelerator.backward(loss)
 
-                        on_before_optimize()
+                        on_before_optimize(epoch)
 
                         optimizer.step()
                         lr_scheduler.step()
diff --git a/training/lr.py b/training/lr.py
index 84e30a0..7584ba2 100644
--- a/training/lr.py
+++ b/training/lr.py
@@ -16,6 +16,10 @@ def noop(*args, **kwards):
     pass
 
 
+def noop_ctx(*args, **kwards):
+    return nullcontext()
+
+
 class LRFinder():
     def __init__(
         self,
@@ -25,10 +29,10 @@ class LRFinder():
         train_dataloader,
         val_dataloader,
         loss_fn: Union[Callable[[int, Any], Tuple[Any, Any, int]], Callable[[int, Any, bool], Tuple[Any, Any, int]]],
-        on_train: Callable[[], _GeneratorContextManager] = nullcontext,
-        on_before_optimize: Callable[[], None] = noop,
+        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] = nullcontext
+        on_eval: Callable[[], _GeneratorContextManager] = noop_ctx
     ):
         self.accelerator = accelerator
         self.model = model
@@ -86,7 +90,7 @@ class LRFinder():
 
             self.model.train()
 
-            with self.on_train():
+            with self.on_train(epoch):
                 for step, batch in enumerate(self.train_dataloader):
                     if step >= num_train_batches:
                         break
@@ -96,7 +100,7 @@ class LRFinder():
 
                         self.accelerator.backward(loss)
 
-                        self.on_before_optimize()
+                        self.on_before_optimize(epoch)
 
                         self.optimizer.step()
                         lr_scheduler.step()
diff --git a/training/modules/dreambooth.py b/training/modules/dreambooth.py
new file mode 100644
index 0000000..e69de29
--- /dev/null
+++ b/training/modules/dreambooth.py
diff --git a/training/modules/lora.py b/training/modules/lora.py
new file mode 100644
index 0000000..e69de29
--- /dev/null
+++ b/training/modules/lora.py
diff --git a/training/modules/ti.py b/training/modules/ti.py
new file mode 100644
index 0000000..2db6f88
--- /dev/null
+++ b/training/modules/ti.py
@@ -0,0 +1,284 @@
+from typing import Literal
+from functools import partial
+from contextlib import contextmanager, nullcontext
+
+import torch
+
+from slugify import slugify
+
+from accelerate import Accelerator
+from transformers import CLIPTextModel
+from diffusers import AutoencoderKL, UNet2DConditionModel
+
+from pipelines.stable_diffusion.vlpn_stable_diffusion import VlpnStableDiffusion
+from models.clip.tokenizer import MultiCLIPTokenizer
+
+from training.common import TrainingSetup, get_scheduler, train_loop, loss_step
+from training.util import EMAModel, CheckpointerBase
+
+
+class Checkpointer(CheckpointerBase):
+    def __init__(
+        self,
+        accelerator: Accelerator,
+        vae: AutoencoderKL,
+        unet: UNet2DConditionModel,
+        tokenizer: MultiCLIPTokenizer,
+        text_encoder: CLIPTextModel,
+        ema_embeddings: EMAModel,
+        weight_dtype: torch.dtype,
+        scheduler,
+        placeholder_token,
+        placeholder_token_ids,
+        *args,
+        **kwargs
+    ):
+        super().__init__(*args, **kwargs)
+
+        self.weight_dtype = weight_dtype
+        self.accelerator = accelerator
+        self.vae = vae
+        self.unet = unet
+        self.tokenizer = tokenizer
+        self.text_encoder = text_encoder
+        self.ema_embeddings = ema_embeddings
+        self.scheduler = scheduler
+        self.placeholder_token = placeholder_token
+        self.placeholder_token_ids = placeholder_token_ids
+
+    @torch.no_grad()
+    def checkpoint(self, step, postfix):
+        print("Saving checkpoint for step %d..." % step)
+
+        checkpoints_path = self.output_dir.joinpath("checkpoints")
+        checkpoints_path.mkdir(parents=True, exist_ok=True)
+
+        text_encoder = self.accelerator.unwrap_model(self.text_encoder)
+
+        ema_context = nullcontext()
+        if self.ema_embeddings is not None:
+            ema_context = self.ema_embeddings.apply_temporary(
+                text_encoder.text_model.embeddings.temp_token_embedding.parameters())
+
+        with ema_context:
+            for (token, ids) in zip(self.placeholder_token, self.placeholder_token_ids):
+                text_encoder.text_model.embeddings.save_embed(
+                    ids,
+                    checkpoints_path.joinpath(f"{slugify(token)}_{step}_{postfix}.bin")
+                )
+
+        del text_encoder
+
+    @torch.no_grad()
+    def save_samples(self, step, num_inference_steps, guidance_scale=7.5, eta=0.0):
+        text_encoder = self.accelerator.unwrap_model(self.text_encoder)
+
+        ema_context = nullcontext()
+        if self.ema_embeddings is not None:
+            ema_context = self.ema_embeddings.apply_temporary(
+                text_encoder.text_model.embeddings.temp_token_embedding.parameters())
+
+        with ema_context:
+            orig_dtype = text_encoder.dtype
+            text_encoder.to(dtype=self.weight_dtype)
+
+            pipeline = VlpnStableDiffusion(
+                text_encoder=text_encoder,
+                vae=self.vae,
+                unet=self.unet,
+                tokenizer=self.tokenizer,
+                scheduler=self.scheduler,
+            ).to(self.accelerator.device)
+            pipeline.set_progress_bar_config(dynamic_ncols=True)
+
+            super().save_samples(pipeline, step, num_inference_steps, guidance_scale, eta)
+
+            text_encoder.to(dtype=orig_dtype)
+
+        del text_encoder
+        del pipeline
+
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+
+
+def train_ti(
+    setup: TrainingSetup,
+    num_train_epochs: int = 100,
+    num_class_images: int = 0,
+    prior_loss_weight: float = 1.0,
+    use_ema: bool = False,
+    ema_inv_gamma: float = 1.0,
+    ema_power: float = 4/5,
+    ema_max_decay: float = .9999,
+    adam_beta1: float = 0.9,
+    adam_beta2: float = 0.999,
+    adam_weight_decay: float = 0,
+    adam_epsilon: float = 1e-08,
+    adam_amsgrad: bool = False,
+    lr_scheduler: Literal[
+        "linear", "cosine", "cosine_with_restarts", "polynomial", "constant", "constant_with_warmup", "one_cycle"
+    ] = "one_cycle",
+    lr_min_lr: float = 0.04,
+    lr_warmup_func: Literal["linear", "cos"] = "cos",
+    lr_annealing_func: Literal["linear", "half_cos", "cos"] = "cos",
+    lr_warmup_exp: int = 1,
+    lr_annealing_exp: int = 1,
+    lr_cycles: int = 1,
+    lr_warmup_epochs: int = 10,
+    emb_decay_target: float = 0.4,
+    emb_decay_factor: float = 1,
+    emb_decay_start: float = 1e-4,
+    sample_image_size: int = 768,
+    sample_batch_size: int = 1,
+    sample_batches: int = 1,
+    sample_frequency: int = 10,
+    sample_steps: int = 20,
+    checkpoint_frequency: int = 50,
+    global_step_offset: int = 0,
+):
+    if use_ema:
+        ema_embeddings = EMAModel(
+            setup.text_encoder.text_model.embeddings.temp_token_embedding.parameters(),
+            inv_gamma=ema_inv_gamma,
+            power=ema_power,
+            max_value=ema_max_decay,
+        )
+    else:
+        ema_embeddings = None
+
+    setup.text_encoder.requires_grad_(True)
+    setup.text_encoder.text_model.encoder.requires_grad_(False)
+    setup.text_encoder.text_model.final_layer_norm.requires_grad_(False)
+    setup.text_encoder.text_model.embeddings.position_embedding.requires_grad_(False)
+    setup.text_encoder.text_model.embeddings.token_embedding.requires_grad_(False)
+
+    # Initialize the optimizer
+    optimizer = setup.optimizer_class(
+        setup.text_encoder.text_model.embeddings.temp_token_embedding.parameters(),
+        lr=setup.learning_rate,
+        betas=(adam_beta1, adam_beta2),
+        weight_decay=adam_weight_decay,
+        eps=adam_epsilon,
+        amsgrad=adam_amsgrad,
+    )
+
+    lr_scheduler = get_scheduler(
+        lr_scheduler,
+        optimizer=optimizer,
+        min_lr=lr_min_lr,
+        warmup_func=lr_warmup_func,
+        annealing_func=lr_annealing_func,
+        warmup_exp=lr_warmup_exp,
+        annealing_exp=lr_annealing_exp,
+        cycles=lr_cycles,
+        train_epochs=num_train_epochs,
+        warmup_epochs=lr_warmup_epochs,
+        num_training_steps_per_epoch=len(setup.train_dataloader),
+        gradient_accumulation_steps=setup.accelerator.gradient_accumulation_steps
+    )
+
+    text_encoder, optimizer, lr_scheduler = setup.accelerator.prepare(
+        setup.text_encoder, optimizer, lr_scheduler
+    )
+
+    # Move vae and unet to device
+    setup.vae.to(setup.accelerator.device, dtype=setup.weight_dtype)
+    setup.unet.to(setup.accelerator.device, dtype=setup.weight_dtype)
+
+    if use_ema:
+        ema_embeddings.to(setup.accelerator.device)
+
+    setup.unet.train()
+
+    @contextmanager
+    def on_train(epoch: int):
+        try:
+            setup.tokenizer.train()
+            yield
+        finally:
+            pass
+
+    @contextmanager
+    def on_eval():
+        try:
+            setup.tokenizer.eval()
+
+            ema_context = nullcontext()
+            if use_ema:
+                ema_context = ema_embeddings.apply_temporary(
+                    text_encoder.text_model.embeddings.temp_token_embedding.parameters())
+
+            with ema_context:
+                yield
+        finally:
+            pass
+
+    @torch.no_grad()
+    def on_after_optimize(lr: float):
+        text_encoder.text_model.embeddings.normalize(
+            emb_decay_target,
+            min(1.0, max(0.0, emb_decay_factor * ((lr - emb_decay_start) / (setup.learning_rate - emb_decay_start))))
+        )
+
+        if use_ema:
+            ema_embeddings.step(text_encoder.text_model.embeddings.temp_token_embedding.parameters())
+
+    def on_log():
+        if use_ema:
+            return {"ema_decay": ema_embeddings.decay}
+        return {}
+
+    loss_step_ = partial(
+        loss_step,
+        setup.vae,
+        setup.noise_scheduler,
+        setup.unet,
+        text_encoder,
+        num_class_images != 0,
+        prior_loss_weight,
+        setup.seed,
+    )
+
+    checkpointer = Checkpointer(
+        accelerator=setup.accelerator,
+        vae=setup.vae,
+        unet=setup.unet,
+        tokenizer=setup.tokenizer,
+        text_encoder=text_encoder,
+        ema_embeddings=ema_embeddings,
+        weight_dtype=setup.weight_dtype,
+        scheduler=setup.checkpoint_scheduler,
+        placeholder_token=setup.placeholder_token,
+        placeholder_token_ids=setup.placeholder_token_ids,
+        train_dataloader=setup.train_dataloader,
+        val_dataloader=setup.val_dataloader,
+        output_dir=setup.output_dir,
+        seed=setup.seed,
+        sample_image_size=sample_image_size,
+        sample_batch_size=sample_batch_size,
+        sample_batches=sample_batches
+    )
+
+    if setup.accelerator.is_main_process:
+        setup.accelerator.init_trackers("textual_inversion")
+
+    train_loop(
+        accelerator=setup.accelerator,
+        optimizer=optimizer,
+        lr_scheduler=lr_scheduler,
+        model=text_encoder,
+        checkpointer=checkpointer,
+        train_dataloader=setup.train_dataloader,
+        val_dataloader=setup.val_dataloader,
+        loss_step=loss_step_,
+        sample_frequency=sample_frequency,
+        sample_steps=sample_steps,
+        checkpoint_frequency=checkpoint_frequency,
+        global_step_offset=global_step_offset,
+        num_epochs=num_train_epochs,
+        on_log=on_log,
+        on_train=on_train,
+        on_after_optimize=on_after_optimize,
+        on_eval=on_eval
+    )
diff --git a/training/util.py b/training/util.py
index 0ec2032..cc4cdee 100644
--- a/training/util.py
+++ b/training/util.py
@@ -41,14 +41,16 @@ class AverageMeter:
 class CheckpointerBase:
     def __init__(
         self,
-        datamodule,
+        train_dataloader,
+        val_dataloader,
         output_dir: Path,
         sample_image_size: int,
         sample_batches: int,
         sample_batch_size: int,
         seed: Optional[int] = None
     ):
-        self.datamodule = datamodule
+        self.train_dataloader = train_dataloader
+        self.val_dataloader = val_dataloader
         self.output_dir = output_dir
         self.sample_image_size = sample_image_size
         self.seed = seed if seed is not None else torch.random.seed()
@@ -70,15 +72,16 @@ class CheckpointerBase:
     ):
         samples_path = Path(self.output_dir).joinpath("samples")
 
-        train_data = self.datamodule.train_dataloader
-        val_data = self.datamodule.val_dataloader
-
         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
 
-        for pool, data, gen in [("stable", val_data, generator), ("val", val_data, None), ("train", train_data, None)]:
+        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)