Made a very basic vae trainer.

.gitmodules

@@ -1,3 +1,6 @@
 [submodule "repositories/sd-scripts"]
 	path = repositories/sd-scripts
 	url = https://github.com/kohya-ss/sd-scripts.git
+[submodule "repositories/leco"]
+	path = repositories/leco
+	url = https://github.com/p1atdev/LECO

jobs/TrainJob.py

@@ -1,85 +1,40 @@
-# from jobs import BaseJob
+from jobs import BaseJob
-# from toolkit.kohya_model_util import load_models_from_stable_diffusion_checkpoint
+from toolkit.kohya_model_util import load_models_from_stable_diffusion_checkpoint
-# from collections import OrderedDict
+from collections import OrderedDict
-# from typing import List
+from typing import List
-# from jobs.process import BaseExtractProcess, TrainFineTuneProcess
+from jobs.process import BaseExtractProcess, TrainFineTuneProcess
-# import gc
+
-# import time
+from toolkit.paths import REPOS_ROOT
-# import argparse
+
-# import itertools
+import sys
-# import math
+
-# import os
+sys.path.append(REPOS_ROOT)
-# from multiprocessing import Value
+
-#
+process_dict = {
-# from tqdm import tqdm
+    'vae': 'TrainVAEProcess',
-# import torch
+    'finetune': 'TrainFineTuneProcess'
-# from accelerate.utils import set_seed
+}
-# from accelerate import Accelerator
+
-# import diffusers
+
-# from diffusers import DDPMScheduler
+class TrainJob(BaseJob):
-#
+    process: List[BaseExtractProcess]
-# from toolkit.paths import SD_SCRIPTS_ROOT
+
-#
+    def __init__(self, config: OrderedDict):
-# import sys
+        super().__init__(config)
-#
+        self.training_folder = self.get_conf('training_folder', required=True)
-# sys.path.append(SD_SCRIPTS_ROOT)
+        self.is_v2 = self.get_conf('is_v2', False)
-#
+        self.device = self.get_conf('device', 'cpu')
-# import library.train_util as train_util
+        self.gradient_accumulation_steps = self.get_conf('gradient_accumulation_steps', 1)
-# import library.config_util as config_util
+        self.mixed_precision = self.get_conf('mixed_precision', False)  # fp16
-# from library.config_util import (
+        self.logging_dir = self.get_conf('logging_dir', None)
-#     ConfigSanitizer,
+
-#     BlueprintGenerator,
+        # loads the processes from the config
-# )
+        self.load_processes(process_dict)
-# import toolkit.train_tools as train_tools
+
-# import library.custom_train_functions as custom_train_functions
+    def run(self):
-# from library.custom_train_functions import (
+        super().run()
-#     apply_snr_weight,
+        print("")
-#     get_weighted_text_embeddings,
+        print(f"Running  {len(self.process)} process{'' if len(self.process) == 1 else 'es'}")
-#     prepare_scheduler_for_custom_training,
+
-#     pyramid_noise_like,
+        for process in self.process:
-#     apply_noise_offset,
+            process.run()
-#     scale_v_prediction_loss_like_noise_prediction,
-# )
-#
-# process_dict = {
-#     'fine_tine': 'TrainFineTuneProcess'
-# }
-#
-#
-# class TrainJob(BaseJob):
-#     process: List[BaseExtractProcess]
-#
-#     def __init__(self, config: OrderedDict):
-#         super().__init__(config)
-#         self.base_model_path = self.get_conf('base_model', required=True)
-#         self.base_model = None
-#         self.training_folder = self.get_conf('training_folder', required=True)
-#         self.is_v2 = self.get_conf('is_v2', False)
-#         self.device = self.get_conf('device', 'cpu')
-#         self.gradient_accumulation_steps = self.get_conf('gradient_accumulation_steps', 1)
-#         self.mixed_precision = self.get_conf('mixed_precision', False)  # fp16
-#         self.logging_dir = self.get_conf('logging_dir', None)
-#
-#         # loads the processes from the config
-#         self.load_processes(process_dict)
-#
-#         # setup accelerator
-#         self.accelerator = Accelerator(
-#             gradient_accumulation_steps=self.gradient_accumulation_steps,
-#             mixed_precision=self.mixed_precision,
-#             log_with=None if self.logging_dir is None else 'tensorboard',
-#             logging_dir=self.logging_dir,
-#         )
-#
-#     def run(self):
-#         super().run()
-#         # load models
-#         print(f"Loading base model for training")
-#         print(f" - Loading base model: {self.base_model_path}")
-#         self.base_model = load_models_from_stable_diffusion_checkpoint(self.is_v2, self.base_model_path)
-#
-#         print("")
-#         print(f"Running  {len(self.process)} process{'' if len(self.process) == 1 else 'es'}")
-#
-#         for process in self.process:
-#             process.run()

jobs/__init__.py

@@ -1,2 +1,3 @@
 from .BaseJob import BaseJob
 from .ExtractJob import ExtractJob
+from .TrainJob import TrainJob

jobs/process/BaseTrainProcess.py

@@ -13,9 +13,6 @@ class BaseTrainProcess(BaseProcess):
             config: OrderedDict
     ):
         super().__init__(process_id, job, config)
-        self.process_id = process_id
-        self.job = job
-        self.config = config
 
     def run(self):
         # implement in child class

jobs/process/TrainVAEProcess.py

@@ -0,0 +1,291 @@
+import copy
+import os
+import time
+from collections import OrderedDict
+
+from PIL import Image
+from PIL.ImageOps import exif_transpose
+from safetensors.torch import save_file
+from torch.utils.data import DataLoader, ConcatDataset
+import torch
+from torch import nn
+from torchvision.transforms import transforms
+
+from jobs.process import BaseTrainProcess
+from toolkit.kohya_model_util import load_vae
+from toolkit.data_loader import ImageDataset
+from toolkit.metadata import get_meta_for_safetensors
+from toolkit.train_tools import get_torch_dtype
+from tqdm import tqdm
+import time
+import numpy as np
+
+IMAGE_TRANSFORMS = transforms.Compose(
+    [
+        transforms.ToTensor(),
+        transforms.Normalize([0.5], [0.5]),
+    ]
+)
+
+INVERSE_IMAGE_TRANSFORMS = transforms.Compose(
+    [
+        transforms.Normalize(
+            mean=[-0.5/0.5],
+            std=[1/0.5]
+        ),
+        transforms.ToPILImage(),
+    ]
+)
+
+
+class TrainVAEProcess(BaseTrainProcess):
+    def __init__(self, process_id: int, job, config: OrderedDict):
+        super().__init__(process_id, job, config)
+        self.data_loader = None
+        self.vae = None
+        self.device = self.get_conf('device', self.job.device)
+        self.vae_path = self.get_conf('vae_path', required=True)
+        self.datasets_objects = self.get_conf('datasets', required=True)
+        self.training_folder = self.get_conf('training_folder', self.job.training_folder)
+        self.batch_size = self.get_conf('batch_size', 1)
+        self.resolution = self.get_conf('resolution', 256)
+        self.learning_rate = self.get_conf('learning_rate', 1e-4)
+        self.sample_every = self.get_conf('sample_every', None)
+        self.epochs = self.get_conf('epochs', None)
+        self.max_steps = self.get_conf('max_steps', None)
+        self.save_every = self.get_conf('save_every', None)
+        self.dtype = self.get_conf('dtype', 'float32')
+        self.sample_sources = self.get_conf('sample_sources', None)
+        self.torch_dtype = get_torch_dtype(self.dtype)
+        self.save_root = os.path.join(self.training_folder, self.job.name)
+
+        if self.sample_every is not None and self.sample_sources is None:
+            raise ValueError('sample_every is specified but sample_sources is not')
+
+        if self.epochs is None and self.max_steps is None:
+            raise ValueError('epochs or max_steps must be specified')
+
+        self.data_loaders = []
+        datasets = []
+        # check datasets
+        assert isinstance(self.datasets_objects, list)
+        for dataset in self.datasets_objects:
+            if 'path' not in dataset:
+                raise ValueError('dataset must have a path')
+            # check if is dir
+            if not os.path.isdir(dataset['path']):
+                raise ValueError(f"dataset path does is not a directory: {dataset['path']}")
+
+        # make training folder
+        if not os.path.exists(self.save_root):
+            os.makedirs(self.save_root, exist_ok=True)
+
+    def load_datasets(self):
+        if self.data_loader is None:
+            print(f"Loading datasets")
+            datasets = []
+            for dataset in self.datasets_objects:
+                print(f" - Dataset: {dataset['path']}")
+                ds = copy.copy(dataset)
+                ds['resolution'] = self.resolution
+                image_dataset = ImageDataset(ds)
+                datasets.append(image_dataset)
+
+            concatenated_dataset = ConcatDataset(datasets)
+            self.data_loader = DataLoader(
+                concatenated_dataset,
+                batch_size=self.batch_size,
+                shuffle=True
+            )
+
+    def get_loss(self, pred, target):
+        loss_fn = nn.MSELoss()
+        loss = loss_fn(pred, target)
+        return loss
+
+    def get_elbo_loss(self, pred, target, mu, log_var):
+        # ELBO (Evidence Lower BOund) loss, aka variational lower bound
+        reconstruction_loss = nn.MSELoss(reduction='sum')
+        BCE = reconstruction_loss(pred, target)  # reconstruction loss
+        KLD = -0.5 * torch.sum(1 + log_var - mu.pow(2) - log_var.exp())  # KL divergence
+        return BCE + KLD
+
+    def save(self, step=None):
+        if not os.path.exists(self.save_root):
+            os.makedirs(self.save_root, exist_ok=True)
+
+        step_num = ''
+        if step is not None:
+            # zeropad 9 digits
+            step_num = f"_{str(step).zfill(9)}"
+
+        filename = f'{self.job.name}{step_num}.safetensors'
+        save_path = os.path.join(self.save_root, filename)
+        # prepare meta
+        save_meta = get_meta_for_safetensors(self.meta, self.job.name)
+
+        state_dict = self.vae.state_dict()
+
+        for key in list(state_dict.keys()):
+            v = state_dict[key]
+            v = v.detach().clone().to("cpu").to(torch.float32)
+            state_dict[key] = v
+
+        # having issues with meta
+        save_file(state_dict, os.path.join(self.save_root, filename), save_meta)
+
+        print(f"Saved to {os.path.join(self.save_root, filename)}")
+
+    def sample(self, step=None):
+        sample_folder = os.path.join(self.save_root, 'samples')
+        if not os.path.exists(sample_folder):
+            os.makedirs(sample_folder, exist_ok=True)
+
+        with torch.no_grad():
+            self.vae.encoder.eval()
+            self.vae.decoder.eval()
+
+            for i, img_url in enumerate(self.sample_sources):
+                img = exif_transpose(Image.open(img_url))
+                img = img.convert('RGB')
+                # crop if not square
+                if img.width != img.height:
+                    min_dim = min(img.width, img.height)
+                    img = img.crop((0, 0, min_dim, min_dim))
+                # resize
+                img = img.resize((self.resolution, self.resolution))
+
+                input_img = img
+                img = IMAGE_TRANSFORMS(img).unsqueeze(0).to(self.device, dtype=self.torch_dtype)
+                decoded = self.vae(img).sample.squeeze(0)
+                decoded = INVERSE_IMAGE_TRANSFORMS(decoded)
+
+                # stack input image and decoded image
+                input_img = input_img.resize((self.resolution, self.resolution))
+                decoded = decoded.resize((self.resolution, self.resolution))
+
+                output_img = Image.new('RGB', (self.resolution * 2, self.resolution))
+                output_img.paste(input_img, (0, 0))
+                output_img.paste(decoded, (self.resolution, 0))
+
+                step_num = ''
+                if step is not None:
+                    # zeropad 9 digits
+                    step_num = f"_{str(step).zfill(9)}"
+                seconds_since_epoch = int(time.time())
+                # zeropad 2 digits
+                i_str = str(i).zfill(2)
+                filename = f"{seconds_since_epoch}{step_num}_{i_str}.png"
+                output_img.save(os.path.join(sample_folder, filename))
+        self.vae.decoder.train()
+
+    def run(self):
+        super().run()
+        self.load_datasets()
+
+        max_step_epochs = self.max_steps // len(self.data_loader)
+        num_epochs = self.epochs
+        if num_epochs is None or num_epochs > max_step_epochs:
+            num_epochs = max_step_epochs
+
+        max_epoch_steps = len(self.data_loader) * num_epochs
+        num_steps = self.max_steps
+        if num_steps is None or num_steps > max_epoch_steps:
+            num_steps = max_epoch_steps
+
+        print(f"Training VAE")
+        print(f" - Training folder: {self.training_folder}")
+        print(f" - Batch size: {self.batch_size}")
+        print(f" - Learning rate: {self.learning_rate}")
+        print(f" - Epochs: {num_epochs}")
+        print(f" - Max steps: {self.max_steps}")
+
+        # load vae
+        print(f"Loading VAE")
+        print(f" - Loading VAE: {self.vae_path}")
+        if self.vae is None:
+            self.vae = load_vae(self.vae_path, dtype=self.torch_dtype)
+
+        # set decoder to train
+        self.vae.to(self.device, dtype=self.torch_dtype)
+        self.vae.requires_grad_(False)
+        self.vae.eval()
+
+        self.vae.decoder.requires_grad_(True)
+        self.vae.decoder.train()
+
+        parameters = self.vae.decoder.parameters()
+
+        optimizer = torch.optim.Adam(parameters, lr=self.learning_rate)
+
+        # setup scheduler
+        # scheduler = lr_scheduler.ConstantLR
+        # todo allow other schedulers
+        scheduler = torch.optim.lr_scheduler.ConstantLR(
+            optimizer,
+            total_iters=num_steps,
+            factor=1,
+            verbose=False
+        )
+
+        # setup tqdm progress bar
+        progress_bar = tqdm(
+            total=num_steps,
+            desc='Training VAE',
+            leave=True
+        )
+
+        step = 0
+        # sample first
+        self.sample()
+        for epoch in range(num_epochs):
+            if step >= num_steps:
+                break
+            for batch in self.data_loader:
+                if step >= num_steps:
+                    break
+
+                batch = batch.to(self.device, dtype=self.torch_dtype)
+
+                # forward pass
+                # with torch.no_grad():
+                dgd = self.vae.encode(batch).latent_dist
+                mu, logvar = dgd.mean, dgd.logvar
+                latents = dgd.sample()
+                latents.requires_grad_(True)
+
+                pred = self.vae.decode(latents).sample
+
+                loss = self.get_elbo_loss(pred, batch, mu, logvar)
+
+                # Backward pass and optimization
+                optimizer.zero_grad()
+                loss.backward()
+                optimizer.step()
+                scheduler.step()
+
+                # update progress bar
+                loss_value = loss.item()
+                # get exponent like 3.54e-4
+                loss_string = f"{loss_value:.2e}"
+                learning_rate = optimizer.param_groups[0]['lr']
+                progress_bar.set_postfix_str(f"LR: {learning_rate:.2e} Loss: {loss_string}")
+                progress_bar.set_description(f"E: {epoch} - S: {step} ")
+                progress_bar.update(1)
+
+                if step != 0:
+                    if self.sample_every and step % self.sample_every == 0:
+                        # print above the progress bar
+                        print(f"Sampling at step {step}")
+                        self.sample(step)
+
+                    if self.save_every and  step % self.save_every == 0:
+                        # print above the progress bar
+                        print(f"Saving at step {step}")
+                        self.save(step)
+
+                step += 1
+
+        self.save()
+
+        pass

jobs/process/__init__.py

@@ -3,4 +3,4 @@ from .ExtractLoconProcess import ExtractLoconProcess
 from .ExtractLoraProcess import ExtractLoraProcess
 from .BaseProcess import BaseProcess
 from .BaseTrainProcess import BaseTrainProcess
-from .TrainFineTuneProcess import TrainFineTuneProcess
+from .TrainVAEProcess import TrainVAEProcess

requirements.txt

@@ -1,4 +1,5 @@
 torch
+torchvision
 safetensors
 diffusers
 transformers

toolkit/data_loader.py

@@ -0,0 +1,67 @@
+import os
+import random
+from PIL import Image
+from PIL.ImageOps import exif_transpose
+from torchvision import transforms
+from torch.utils.data import Dataset
+
+
+class ImageDataset(Dataset):
+    def __init__(self, config):
+        self.config = config
+        self.name = self.get_config('name', 'dataset')
+        self.path = self.get_config('path', required=True)
+        self.scale = self.get_config('scale', 1)
+        self.random_scale = self.get_config('random_scale', False)
+        # we always random crop if random scale is enabled
+        self.random_crop = self.random_scale if self.random_scale else self.get_config('random_crop', False)
+
+        self.resolution = self.get_config('resolution', 256)
+        self.file_list = [os.path.join(self.path, file) for file in os.listdir(self.path) if
+                          file.lower().endswith(('.jpg', '.jpeg', '.png', '.webp'))]
+
+        # this might take a while
+        print(f"  -  Preprocessing image dimensions")
+        self.file_list = [file for file in self.file_list if
+                          int(min(Image.open(file).size) * self.scale) >= self.resolution]
+
+        print(f"  -  Found {len(self.file_list)} images")
+        assert len(self.file_list) > 0, f"no images found in {self.path}"
+
+        self.transform = transforms.Compose([
+            transforms.ToTensor(),
+            transforms.Normalize([0.5], [0.5]),
+        ])
+
+    def get_config(self, key, default=None, required=False):
+        if key in self.config:
+            value = self.config[key]
+            return value
+        elif required:
+            raise ValueError(f'config file error. Missing "config.dataset.{key}" key')
+        else:
+            return default
+
+    def __len__(self):
+        return len(self.file_list)
+
+    def __getitem__(self, index):
+        img_path = self.file_list[index]
+        img = exif_transpose(Image.open(img_path)).convert('RGB')
+
+        # Downscale the source image first
+        img = img.resize((int(img.size[0] * self.scale), int(img.size[1] * self.scale)), Image.BICUBIC)
+
+        if self.random_crop:
+            if self.random_scale:
+                scale_size = random.randint(int(img.size[0] * self.scale), self.resolution)
+                img = img.resize((scale_size, scale_size), Image.BICUBIC)
+            img = transforms.RandomCrop(self.resolution)(img)
+        else:
+            min_dimension = min(img.size)
+            img = transforms.CenterCrop(min_dimension)(img)
+            img = img.resize((self.resolution, self.resolution), Image.BICUBIC)
+
+        img = self.transform(img)
+
+        return img

toolkit/job.py

@@ -10,6 +10,10 @@ def get_job(config_path):
     if job == 'extract':
         from jobs import ExtractJob
         return ExtractJob(config)
+    if job == 'train':
+        from jobs import TrainJob
+        return TrainJob(config)
+
     # elif job == 'train':
     #     from jobs import TrainJob
     #     return TrainJob(config)

toolkit/paths.py

@@ -3,3 +3,4 @@ import os
 TOOLKIT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
 CONFIG_ROOT = os.path.join(TOOLKIT_ROOT, 'config')
 SD_SCRIPTS_ROOT = os.path.join(TOOLKIT_ROOT, "repositories", "sd-scripts")
+REPOS_ROOT = os.path.join(TOOLKIT_ROOT, "repositories")