6 files changed, 33 insertions, 127 deletions

diff --git a/train_ti.py b/train_ti.py
index a4e2dde..78c1b5c 100644
--- a/train_ti.py
+++ b/train_ti.py
@@ -11,20 +11,16 @@ import torch.utils.checkpoint
 from accelerate import Accelerator
 from accelerate.logging import get_logger
 from accelerate.utils import LoggerType, set_seed
-from diffusers import AutoencoderKL, UNet2DConditionModel
 import matplotlib.pyplot as plt
-from transformers import CLIPTextModel
 from slugify import slugify
 
 from util import load_config, load_embeddings_from_dir
-from pipelines.stable_diffusion.vlpn_stable_diffusion import VlpnStableDiffusion
 from data.csv import VlpnDataModule, VlpnDataItem
-from trainer.base import Checkpointer
+from trainer_old.base import Checkpointer
 from training.functional import loss_step, train_loop, generate_class_images, add_placeholder_tokens, get_models
 from training.optimization import get_scheduler
 from training.lr import LRFinder
 from training.util import EMAModel, save_args
-from models.clip.tokenizer import MultiCLIPTokenizer
 
 logger = get_logger(__name__)
 
@@ -485,12 +481,16 @@ class TextualInversionCheckpointer(Checkpointer):
     def __init__(
         self,
         ema_embeddings: EMAModel,
+        placeholder_tokens: list[str],
+        placeholder_token_ids: list[list[int]],
         *args,
         **kwargs,
     ):
         super().__init__(*args, **kwargs)
 
         self.ema_embeddings = ema_embeddings
+        self.placeholder_tokens = placeholder_tokens
+        self.placeholder_token_ids = placeholder_token_ids
 
     @torch.no_grad()
     def checkpoint(self, step, postfix):
diff --git a/trainer/base.py b/trainer_old/base.py
index 1f85e71..1f85e71 100644
--- a/trainer/base.py
+++ b/trainer_old/base.py
diff --git a/trainer/dreambooth.py b/trainer_old/dreambooth.py
index e69de29..e69de29 100644
--- a/trainer/dreambooth.py
+++ b/trainer_old/dreambooth.py
diff --git a/trainer/ti.py b/trainer_old/ti.py
index 388acd3..66393af 100644
--- a/trainer/ti.py
+++ b/trainer_old/ti.py
@@ -15,12 +15,16 @@ class TextualInversionCheckpointer(Checkpointer):
     def __init__(
         self,
         ema_embeddings: EMAModel,
+        placeholder_tokens: list[str],
+        placeholder_token_ids: list[list[int]],
         *args,
         **kwargs,
     ):
         super().__init__(*args, **kwargs)
 
         self.ema_embeddings = ema_embeddings
+        self.placeholder_tokens = placeholder_tokens
+        self.placeholder_token_ids = placeholder_token_ids
 
     @torch.no_grad()
     def checkpoint(self, step, postfix):
diff --git a/training/functional.py b/training/functional.py
index c100ea2..c5b514a 100644
--- a/training/functional.py
+++ b/training/functional.py
@@ -25,17 +25,31 @@ def const(result=None):
     return fn
 
 
+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')
+
+    embeddings = patch_managed_embeddings(text_encoder)
+
+    return tokenizer, text_encoder, vae, unet, noise_scheduler, sample_scheduler, embeddings
+
+
 def generate_class_images(
-    accelerator,
-    text_encoder,
-    vae,
-    unet,
-    tokenizer,
-    scheduler,
+    accelerator: Accelerator,
+    text_encoder: CLIPTextModel,
+    vae: AutoencoderKL,
+    unet: UNet2DConditionModel,
+    tokenizer: MultiCLIPTokenizer,
+    sample_scheduler: DPMSolverMultistepScheduler,
     data_train,
-    sample_batch_size,
-    sample_image_size,
-    sample_steps
+    sample_batch_size: int,
+    sample_image_size: int,
+    sample_steps: int
 ):
     missing_data = [item for item in data_train if not item.class_image_path.exists()]
 
@@ -52,7 +66,7 @@ def generate_class_images(
         vae=vae,
         unet=unet,
         tokenizer=tokenizer,
-        scheduler=scheduler,
+        scheduler=sample_scheduler,
     ).to(accelerator.device)
     pipeline.set_progress_bar_config(dynamic_ncols=True)
 
diff --git a/training/util.py b/training/util.py
index a292edd..f46cc61 100644
--- a/training/util.py
+++ b/training/util.py
@@ -14,29 +14,6 @@ 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={}):
     info = {"args": vars(args)}
     info["args"].update(extra)
@@ -44,95 +21,6 @@ def save_args(basepath: Path, args, extra={}):
         json.dump(info, f, indent=4)
 
 
-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
-
-    batched_data = [
-        missing_data[i:i+sample_batch_size]
-        for i in range(0, len(missing_data), sample_batch_size)
-    ]
-
-    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)
-
-    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
-
-            for i, image in enumerate(images):
-                image.save(image_name[i])
-
-    del pipeline
-
-    if torch.cuda.is_available():
-        torch.cuda.empty_cache()
-
-
-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')
-
-    embeddings = patch_managed_embeddings(text_encoder)
-
-    return tokenizer, text_encoder, vae, unet, noise_scheduler, sample_scheduler, embeddings
-
-
-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)
-
-    embeddings.resize(len(tokenizer))
-
-    for (placeholder_token_id, initializer_token_id) in zip(placeholder_token_ids, initializer_token_ids):
-        embeddings.add_embed(placeholder_token_id, initializer_token_id)
-
-    return placeholder_token_ids, initializer_token_ids
-
-
 class AverageMeter:
     def __init__(self, name=None):
         self.name = name