Moved the base SD training process to a parent class so we can do some more creative stuff without copy pasta

jobs/process/BaseSDTrainProcess.py

@@ -0,0 +1,420 @@
+import time
+from collections import OrderedDict
+import os
+
+from toolkit.kohya_model_util import load_vae
+from toolkit.lora_special import LoRASpecialNetwork
+from toolkit.optimizer import get_optimizer
+from toolkit.paths import REPOS_ROOT
+import sys
+
+sys.path.append(REPOS_ROOT)
+sys.path.append(os.path.join(REPOS_ROOT, 'leco'))
+
+from diffusers import StableDiffusionPipeline
+
+from jobs.process import BaseTrainProcess
+from toolkit.metadata import get_meta_for_safetensors
+from toolkit.train_tools import get_torch_dtype, apply_noise_offset
+import gc
+
+import torch
+from tqdm import tqdm
+
+from leco import train_util, model_util
+from toolkit.config_modules import SaveConfig, LogingConfig, SampleConfig, NetworkConfig, TrainConfig, ModelConfig
+
+
+def flush():
+    torch.cuda.empty_cache()
+    gc.collect()
+
+
+UNET_IN_CHANNELS = 4  # Stable Diffusion の in_channels は 4 で固定。XLも同じ。
+VAE_SCALE_FACTOR = 8  # 2 ** (len(vae.config.block_out_channels) - 1) = 8
+
+
+class StableDiffusion:
+    def __init__(self, vae, tokenizer, text_encoder, unet, noise_scheduler):
+        self.vae = vae
+        self.tokenizer = tokenizer
+        self.text_encoder = text_encoder
+        self.unet = unet
+        self.noise_scheduler = noise_scheduler
+
+
+class BaseSDTrainProcess(BaseTrainProcess):
+    def __init__(self, process_id: int, job, config: OrderedDict):
+        super().__init__(process_id, job, config)
+        self.step_num = 0
+        self.start_step = 0
+        self.device = self.get_conf('device', self.job.device)
+        self.device_torch = torch.device(self.device)
+        self.network_config = NetworkConfig(**self.get_conf('network', None))
+        self.training_folder = self.get_conf('training_folder', self.job.training_folder)
+        self.train_config = TrainConfig(**self.get_conf('train', {}))
+        self.model_config = ModelConfig(**self.get_conf('model', {}))
+        self.save_config = SaveConfig(**self.get_conf('save', {}))
+        self.sample_config = SampleConfig(**self.get_conf('sample', {}))
+        self.logging_config = LogingConfig(**self.get_conf('logging', {}))
+        self.optimizer = None
+        self.lr_scheduler = None
+        self.sd = None
+
+        # added later
+        self.network = None
+        self.scheduler = None
+
+    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)
+
+        if self.network is not None:
+            self.network.eval()
+
+        # save current seed state for training
+        rng_state = torch.get_rng_state()
+        cuda_rng_state = torch.cuda.get_rng_state() if torch.cuda.is_available() else None
+
+        original_device_dict = {
+            'vae': self.sd.vae.device,
+            'unet': self.sd.unet.device,
+            'text_encoder': self.sd.text_encoder.device,
+            # 'tokenizer': self.sd.tokenizer.device,
+        }
+
+        self.sd.vae.to(self.device_torch)
+        self.sd.unet.to(self.device_torch)
+        self.sd.text_encoder.to(self.device_torch)
+        # self.sd.tokenizer.to(self.device_torch)
+        # TODO add clip skip
+
+        pipeline = StableDiffusionPipeline(
+            vae=self.sd.vae,
+            unet=self.sd.unet,
+            text_encoder=self.sd.text_encoder,
+            tokenizer=self.sd.tokenizer,
+            scheduler=self.sd.noise_scheduler,
+            safety_checker=None,
+            feature_extractor=None,
+            requires_safety_checker=False,
+        )
+        # disable progress bar
+        pipeline.set_progress_bar_config(disable=True)
+
+        start_seed = self.sample_config.seed
+        start_multiplier = self.network.multiplier
+        current_seed = start_seed
+
+        pipeline.to(self.device_torch)
+        with self.network:
+            with torch.no_grad():
+                if self.network is not None:
+                    assert self.network.is_active
+                    if self.logging_config.verbose:
+                        print("network_state", {
+                            'is_active': self.network.is_active,
+                            'multiplier': self.network.multiplier,
+                        })
+
+                for i in tqdm(range(len(self.sample_config.prompts)), desc=f"Generating Samples - step: {step}"):
+                    raw_prompt = self.sample_config.prompts[i]
+
+                    neg = self.sample_config.neg
+                    multiplier = self.sample_config.network_multiplier
+                    p_split = raw_prompt.split('--')
+                    prompt = p_split[0].strip()
+
+                    if len(p_split) > 1:
+                        for split in p_split:
+                            flag = split[:1]
+                            content = split[1:].strip()
+                            if flag == 'n':
+                                neg = content
+                            elif flag == 'm':
+                                # multiplier
+                                multiplier = float(content)
+
+                    height = self.sample_config.height
+                    width = self.sample_config.width
+                    height = max(64, height - height % 8)  # round to divisible by 8
+                    width = max(64, width - width % 8)  # round to divisible by 8
+
+                    if self.sample_config.walk_seed:
+                        current_seed += i
+
+                    if self.network is not None:
+                        self.network.multiplier = multiplier
+                    torch.manual_seed(current_seed)
+                    torch.cuda.manual_seed(current_seed)
+
+                    img = pipeline(
+                        prompt,
+                        height=height,
+                        width=width,
+                        num_inference_steps=self.sample_config.sample_steps,
+                        guidance_scale=self.sample_config.guidance_scale,
+                        negative_prompt=neg,
+                    ).images[0]
+
+                    step_num = ''
+                    if step is not None:
+                        # zero-pad 9 digits
+                        step_num = f"_{str(step).zfill(9)}"
+                    seconds_since_epoch = int(time.time())
+                    # zero-pad 2 digits
+                    i_str = str(i).zfill(2)
+                    filename = f"{seconds_since_epoch}{step_num}_{i_str}.png"
+                    output_path = os.path.join(sample_folder, filename)
+                    img.save(output_path)
+
+        # clear pipeline and cache to reduce vram usage
+        del pipeline
+        torch.cuda.empty_cache()
+
+        # restore training state
+        torch.set_rng_state(rng_state)
+        if cuda_rng_state is not None:
+            torch.cuda.set_rng_state(cuda_rng_state)
+
+        self.sd.vae.to(original_device_dict['vae'])
+        self.sd.unet.to(original_device_dict['unet'])
+        self.sd.text_encoder.to(original_device_dict['text_encoder'])
+        if self.network is not None:
+            self.network.train()
+            self.network.multiplier = start_multiplier
+        # self.sd.tokenizer.to(original_device_dict['tokenizer'])
+
+    def update_training_metadata(self):
+        self.add_meta(OrderedDict({"training_info": self.get_training_info()}))
+
+    def get_training_info(self):
+        info = OrderedDict({
+            'step': self.step_num + 1
+        })
+        return info
+
+    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)}"
+
+        self.update_training_metadata()
+        filename = f'{self.job.name}{step_num}.safetensors'
+        file_path = os.path.join(self.save_root, filename)
+        # prepare meta
+        save_meta = get_meta_for_safetensors(self.meta, self.job.name)
+        if self.network is not None:
+            # TODO handle dreambooth, fine tuning, etc
+            self.network.save_weights(
+                file_path,
+                dtype=get_torch_dtype(self.save_config.dtype),
+                metadata=save_meta
+            )
+        else:
+            # TODO handle dreambooth, fine tuning, etc
+            # will probably have to convert dict back to LDM
+            ValueError("Non network training is not currently supported")
+
+        self.print(f"Saved to {file_path}")
+
+    # Called before the model is loaded
+    def hook_before_model_load(self):
+        # override in subclass
+        pass
+
+    def hook_add_extra_train_params(self, params):
+        # override in subclass
+        return params
+
+    def hook_before_train_loop(self):
+        pass
+
+    def get_latent_noise(
+            self,
+            height=None,
+            width=None,
+            pixel_height=None,
+            pixel_width=None,
+    ):
+        if height is None and pixel_height is None:
+            raise ValueError("height or pixel_height must be specified")
+        if width is None and pixel_width is None:
+            raise ValueError("width or pixel_width must be specified")
+        if height is None:
+            height = pixel_height // VAE_SCALE_FACTOR
+        if width is None:
+            width = pixel_width // VAE_SCALE_FACTOR
+
+        noise = torch.randn(
+            (
+                self.train_config.batch_size,
+                UNET_IN_CHANNELS,
+                height,
+                width,
+            ),
+            device="cpu",
+        )
+        noise = apply_noise_offset(noise, self.train_config.noise_offset)
+        return noise
+
+    def hook_train_loop(self):
+        # return loss
+        return 0.0
+
+    def run(self):
+        super().run()
+
+        ### HOOK ###
+        self.hook_before_model_load()
+
+        dtype = get_torch_dtype(self.train_config.dtype)
+
+        tokenizer, text_encoder, unet, noise_scheduler = model_util.load_models(
+            self.model_config.name_or_path,
+            scheduler_name=self.train_config.noise_scheduler,
+            v2=self.model_config.is_v2,
+            v_pred=self.model_config.is_v_pred,
+        )
+        # just for now or of we want to load a custom one
+        # put on cpu for now, we only need it when sampling
+        vae = load_vae(self.model_config.name_or_path, dtype=dtype).to('cpu', dtype=dtype)
+        vae.eval()
+        self.sd = StableDiffusion(vae, tokenizer, text_encoder, unet, noise_scheduler)
+
+        text_encoder.to(self.device_torch, dtype=dtype)
+        text_encoder.eval()
+
+        unet.to(self.device_torch, dtype=dtype)
+        if self.train_config.xformers:
+            unet.enable_xformers_memory_efficient_attention()
+        unet.requires_grad_(False)
+        unet.eval()
+
+        if self.network_config is not None:
+            self.network = LoRASpecialNetwork(
+                text_encoder=text_encoder,
+                unet=unet,
+                lora_dim=self.network_config.rank,
+                multiplier=1.0,
+                alpha=self.network_config.alpha
+            )
+
+            self.network.force_to(self.device_torch, dtype=dtype)
+
+            self.network.apply_to(
+                text_encoder,
+                unet,
+                self.train_config.train_text_encoder,
+                self.train_config.train_unet
+            )
+
+            self.network.prepare_grad_etc(text_encoder, unet)
+
+            params = self.network.prepare_optimizer_params(
+                text_encoder_lr=self.train_config.lr,
+                unet_lr=self.train_config.lr,
+                default_lr=self.train_config.lr
+            )
+
+        else:
+            params = []
+            # assume dreambooth/finetune
+            if self.train_config.train_text_encoder:
+                text_encoder.requires_grad_(True)
+                text_encoder.train()
+                params += text_encoder.parameters()
+            if self.train_config.train_unet:
+                unet.requires_grad_(True)
+                unet.train()
+                params += unet.parameters()
+
+        ### HOOK ###
+        params = self.hook_add_extra_train_params(params)
+
+        optimizer_type = self.train_config.optimizer.lower()
+        optimizer = get_optimizer(params, optimizer_type, learning_rate=self.train_config.lr,
+                                  optimizer_params=self.train_config.optimizer_params)
+        self.optimizer = optimizer
+
+        lr_scheduler = train_util.get_lr_scheduler(
+            self.train_config.lr_scheduler,
+            optimizer,
+            max_iterations=self.train_config.steps,
+            lr_min=self.train_config.lr / 100,  # not sure why leco did this, but ill do it to
+        )
+
+        self.lr_scheduler = lr_scheduler
+
+        ### HOOK ###
+        self.hook_before_train_loop()
+
+        # sample first
+        self.print("Generating baseline samples before training")
+        self.sample(0)
+
+        self.progress_bar = tqdm(
+            total=self.train_config.steps,
+            desc=self.job.name,
+            leave=True
+        )
+        self.step_num = 0
+        for step in range(self.train_config.steps):
+            # todo handle dataloader here maybe, not sure
+
+            ### HOOK ###
+            loss = self.hook_train_loop()
+
+            # don't do on first step
+            if self.step_num != self.start_step:
+                # pause progress bar
+                self.progress_bar.unpause()  # makes it so doesn't track time
+                if self.sample_config.sample_every and self.step_num % self.sample_config.sample_every == 0:
+                    # print above the progress bar
+                    self.sample(self.step_num)
+
+                if self.save_config.save_every and self.step_num % self.save_config.save_every == 0:
+                    # print above the progress bar
+                    self.print(f"Saving at step {self.step_num}")
+                    self.save(self.step_num)
+
+                if self.logging_config.log_every and self.step_num % self.logging_config.log_every == 0:
+                    # log to tensorboard
+                    if self.writer is not None:
+                        # get avg loss
+                        self.writer.add_scalar(f"loss", loss, self.step_num)
+                        if self.train_config.optimizer.startswith('dadaptation'):
+                            learning_rate = (
+                                    optimizer.param_groups[0]["d"] *
+                                    optimizer.param_groups[0]["lr"]
+                            )
+                        else:
+                            learning_rate = optimizer.param_groups[0]['lr']
+                        self.writer.add_scalar(f"lr", learning_rate, self.step_num)
+                self.progress_bar.refresh()
+
+            # sets progress bar to match out step
+            self.progress_bar.update(step - self.progress_bar.n)
+            # end of step
+            self.step_num = step
+
+        self.sample(self.step_num + 1)
+        print("")
+        self.save()
+
+        del (
+            self.sd,
+            unet,
+            noise_scheduler,
+            optimizer,
+            self.network,
+            tokenizer,
+            text_encoder,
+        )
+
+        flush()

jobs/process/TrainSliderProcess.py

@@ -3,30 +3,20 @@
 import time
 from collections import OrderedDict
 import os
-from typing import List, Literal
 
-from toolkit.kohya_model_util import load_vae
+from toolkit.config_modules import SliderConfig
-from toolkit.lora_special import LoRASpecialNetwork
-from toolkit.optimizer import get_optimizer
 from toolkit.paths import REPOS_ROOT
 import sys
 
 sys.path.append(REPOS_ROOT)
 sys.path.append(os.path.join(REPOS_ROOT, 'leco'))
-
-from diffusers import StableDiffusionPipeline
-
-from jobs.process import BaseTrainProcess
-from toolkit.metadata import get_meta_for_safetensors
 from toolkit.train_tools import get_torch_dtype, apply_noise_offset
 import gc
 
 import torch
-from tqdm import tqdm
-
-from toolkit.lora import LoRANetwork, DEFAULT_TARGET_REPLACE, UNET_TARGET_REPLACE_MODULE_CONV, TRAINING_METHODS
 from leco import train_util, model_util
 from leco.prompt_util import PromptEmbedsCache
+from .BaseSDTrainProcess import BaseSDTrainProcess, StableDiffusion
 
 
 class ACTION_TYPES_SLIDER:
@@ -39,97 +29,6 @@ def flush():
     gc.collect()
 
 
-UNET_IN_CHANNELS = 4  # Stable Diffusion の in_channels は 4 で固定。XLも同じ。
-VAE_SCALE_FACTOR = 8  # 2 ** (len(vae.config.block_out_channels) - 1) = 8
-
-
-class StableDiffusion:
-    def __init__(self, vae, tokenizer, text_encoder, unet, noise_scheduler):
-        self.vae = vae
-        self.tokenizer = tokenizer
-        self.text_encoder = text_encoder
-        self.unet = unet
-        self.noise_scheduler = noise_scheduler
-
-
-class SaveConfig:
-    def __init__(self, **kwargs):
-        self.save_every: int = kwargs.get('save_every', 1000)
-        self.dtype: str = kwargs.get('save_dtype', 'float16')
-
-
-class LogingConfig:
-    def __init__(self, **kwargs):
-        self.log_every: int = kwargs.get('log_every', 100)
-        self.verbose: bool = kwargs.get('verbose', False)
-        self.use_wandb: bool = kwargs.get('use_wandb', False)
-
-
-class SampleConfig:
-    def __init__(self, **kwargs):
-        self.sample_every: int = kwargs.get('sample_every', 100)
-        self.width: int = kwargs.get('width', 512)
-        self.height: int = kwargs.get('height', 512)
-        self.prompts: list[str] = kwargs.get('prompts', [])
-        self.neg = kwargs.get('neg', False)
-        self.seed = kwargs.get('seed', 0)
-        self.walk_seed = kwargs.get('walk_seed', False)
-        self.guidance_scale = kwargs.get('guidance_scale', 7)
-        self.sample_steps = kwargs.get('sample_steps', 20)
-        self.network_multiplier = kwargs.get('network_multiplier', 1)
-
-
-class NetworkConfig:
-    def __init__(self, **kwargs):
-        self.type: str = kwargs.get('type', 'lierla')
-        self.rank: int = kwargs.get('rank', 4)
-        self.alpha: float = kwargs.get('alpha', 1.0)
-
-
-class TrainConfig:
-    def __init__(self, **kwargs):
-        self.noise_scheduler: 'model_util.AVAILABLE_SCHEDULERS' = kwargs.get('noise_scheduler', 'ddpm')
-        self.steps: int = kwargs.get('steps', 1000)
-        self.lr = kwargs.get('lr', 1e-6)
-        self.optimizer = kwargs.get('optimizer', 'adamw')
-        self.lr_scheduler = kwargs.get('lr_scheduler', 'constant')
-        self.max_denoising_steps: int = kwargs.get('max_denoising_steps', 50)
-        self.batch_size: int = kwargs.get('batch_size', 1)
-        self.dtype: str = kwargs.get('dtype', 'fp32')
-        self.xformers = kwargs.get('xformers', False)
-        self.train_unet = kwargs.get('train_unet', True)
-        self.train_text_encoder = kwargs.get('train_text_encoder', True)
-        self.noise_offset = kwargs.get('noise_offset', 0.0)
-        self.optimizer_params = kwargs.get('optimizer_params', {})
-
-
-class ModelConfig:
-    def __init__(self, **kwargs):
-        self.name_or_path: str = kwargs.get('name_or_path', None)
-        self.is_v2: bool = kwargs.get('is_v2', False)
-        self.is_v_pred: bool = kwargs.get('is_v_pred', False)
-
-        if self.name_or_path is None:
-            raise ValueError('name_or_path must be specified')
-
-
-class SliderTargetConfig:
-    def __init__(self, **kwargs):
-        self.target_class: str = kwargs.get('target_class', '')
-        self.positive: str = kwargs.get('positive', None)
-        self.negative: str = kwargs.get('negative', None)
-        self.multiplier: float = kwargs.get('multiplier', 1.0)
-        self.weight: float = kwargs.get('weight', 1.0)
-
-
-class SliderConfig:
-    def __init__(self, **kwargs):
-        targets = kwargs.get('targets', [])
-        targets = [SliderTargetConfig(**target) for target in targets]
-        self.targets: List[SliderTargetConfig] = targets
-        self.resolutions: List[List[int]] = kwargs.get('resolutions', [[512, 512]])
-
-
 class EncodedPromptPair:
     def __init__(
             self,
@@ -154,248 +53,22 @@ class EncodedPromptPair:
         self.weight = weight
 
 
-class TrainSliderProcess(BaseTrainProcess):
+class TrainSliderProcess(BaseSDTrainProcess):
     def __init__(self, process_id: int, job, config: OrderedDict):
         super().__init__(process_id, job, config)
         self.step_num = 0
         self.start_step = 0
         self.device = self.get_conf('device', self.job.device)
         self.device_torch = torch.device(self.device)
-        self.network_config = NetworkConfig(**self.get_conf('network', {}))
-        self.training_folder = self.get_conf('training_folder', self.job.training_folder)
-        self.train_config = TrainConfig(**self.get_conf('train', {}))
-        self.model_config = ModelConfig(**self.get_conf('model', {}))
-        self.save_config = SaveConfig(**self.get_conf('save', {}))
-        self.sample_config = SampleConfig(**self.get_conf('sample', {}))
-        self.logging_config = LogingConfig(**self.get_conf('logging', {}))
         self.slider_config = SliderConfig(**self.get_conf('slider', {}))
-        self.sd = None
 
-        # added later
+        self.prompt_cache = PromptEmbedsCache()
-        self.network = None
+        self.prompt_pairs: list[EncodedPromptPair] = []
-        self.scheduler = None
-        self.is_flipped = False
 
-    def flip_network(self):
+    def before_model_load(self):
-        for param in self.network.parameters():
+        pass
-            # apply opposite weight to the network
-            param.data = -param.data
-        self.is_flipped = not self.is_flipped
-
-    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)
-
-        self.network.eval()
-
-        # save current seed state for training
-        rng_state = torch.get_rng_state()
-        cuda_rng_state = torch.cuda.get_rng_state() if torch.cuda.is_available() else None
-
-        original_device_dict = {
-            'vae': self.sd.vae.device,
-            'unet': self.sd.unet.device,
-            'text_encoder': self.sd.text_encoder.device,
-            # 'tokenizer': self.sd.tokenizer.device,
-        }
-
-        self.sd.vae.to(self.device_torch)
-        self.sd.unet.to(self.device_torch)
-        self.sd.text_encoder.to(self.device_torch)
-        # self.sd.tokenizer.to(self.device_torch)
-        # TODO add clip skip
-
-        pipeline = StableDiffusionPipeline(
-            vae=self.sd.vae,
-            unet=self.sd.unet,
-            text_encoder=self.sd.text_encoder,
-            tokenizer=self.sd.tokenizer,
-            scheduler=self.sd.noise_scheduler,
-            safety_checker=None,
-            feature_extractor=None,
-            requires_safety_checker=False,
-        )
-        # disable progress bar
-        pipeline.set_progress_bar_config(disable=True)
-
-        start_seed = self.sample_config.seed
-        start_multiplier = self.network.multiplier
-        current_seed = start_seed
-
-        pipeline.to(self.device_torch)
-        with self.network:
-            with torch.no_grad():
-                assert self.network.is_active
-                if self.logging_config.verbose:
-                    print("network_state", {
-                        'is_active': self.network.is_active,
-                        'multiplier': self.network.multiplier,
-                    })
-
-                for i in tqdm(range(len(self.sample_config.prompts)), desc=f"Generating Samples - step: {step}"):
-                    raw_prompt = self.sample_config.prompts[i]
-
-                    neg = self.sample_config.neg
-                    multiplier = self.sample_config.network_multiplier
-                    p_split = raw_prompt.split('--')
-                    prompt = p_split[0].strip()
-
-                    if len(p_split) > 1:
-                        for split in p_split:
-                            flag = split[:1]
-                            content = split[1:].strip()
-                            if flag == 'n':
-                                neg = content
-                            elif flag == 'm':
-                                # multiplier
-                                multiplier = float(content)
-
-                    height = self.sample_config.height
-                    width = self.sample_config.width
-                    height = max(64, height - height % 8)  # round to divisible by 8
-                    width = max(64, width - width % 8)  # round to divisible by 8
-
-                    if self.sample_config.walk_seed:
-                        current_seed += i
-
-                    self.network.multiplier = multiplier
-                    torch.manual_seed(current_seed)
-                    torch.cuda.manual_seed(current_seed)
-
-                    img = pipeline(
-                        prompt,
-                        height=height,
-                        width=width,
-                        num_inference_steps=self.sample_config.sample_steps,
-                        guidance_scale=self.sample_config.guidance_scale,
-                        negative_prompt=neg,
-                    ).images[0]
-
-                    step_num = ''
-                    if step is not None:
-                        # zero-pad 9 digits
-                        step_num = f"_{str(step).zfill(9)}"
-                    seconds_since_epoch = int(time.time())
-                    # zero-pad 2 digits
-                    i_str = str(i).zfill(2)
-                    filename = f"{seconds_since_epoch}{step_num}_{i_str}.png"
-                    output_path = os.path.join(sample_folder, filename)
-                    img.save(output_path)
-
-        # clear pipeline and cache to reduce vram usage
-        del pipeline
-        torch.cuda.empty_cache()
-
-        # restore training state
-        torch.set_rng_state(rng_state)
-        if cuda_rng_state is not None:
-            torch.cuda.set_rng_state(cuda_rng_state)
-
-        self.sd.vae.to(original_device_dict['vae'])
-        self.sd.unet.to(original_device_dict['unet'])
-        self.sd.text_encoder.to(original_device_dict['text_encoder'])
-        self.network.train()
-        self.network.multiplier = start_multiplier
-        # self.sd.tokenizer.to(original_device_dict['tokenizer'])
-
-    def update_training_metadata(self):
-        self.add_meta(OrderedDict({"training_info": self.get_training_info()}))
-
-    def get_training_info(self):
-        info = OrderedDict({
-            'step': self.step_num + 1
-        })
-        return info
-
-    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)}"
-
-        self.update_training_metadata()
-        filename = f'{self.job.name}{step_num}.safetensors'
-        file_path = os.path.join(self.save_root, filename)
-        # prepare meta
-        save_meta = get_meta_for_safetensors(self.meta, self.job.name)
-        self.network.save_weights(
-            file_path,
-            dtype=get_torch_dtype(self.save_config.dtype),
-            metadata=save_meta
-        )
-
-        self.print(f"Saved to {file_path}")
-
-    def run(self):
-        super().run()
-
-        dtype = get_torch_dtype(self.train_config.dtype)
-
-        modules = DEFAULT_TARGET_REPLACE
-        loss = None
-        if self.network_config.type == "c3lier":
-            modules += UNET_TARGET_REPLACE_MODULE_CONV
-
-        tokenizer, text_encoder, unet, noise_scheduler = model_util.load_models(
-            self.model_config.name_or_path,
-            scheduler_name=self.train_config.noise_scheduler,
-            v2=self.model_config.is_v2,
-            v_pred=self.model_config.is_v_pred,
-        )
-        # just for now or of we want to load a custom one
-        # put on cpu for now, we only need it when sampling
-        vae = load_vae(self.model_config.name_or_path, dtype=dtype).to('cpu', dtype=dtype)
-        vae.eval()
-        self.sd = StableDiffusion(vae, tokenizer, text_encoder, unet, noise_scheduler)
-
-        text_encoder.to(self.device_torch, dtype=dtype)
-        text_encoder.eval()
-
-        unet.to(self.device_torch, dtype=dtype)
-        if self.train_config.xformers:
-            unet.enable_xformers_memory_efficient_attention()
-        unet.requires_grad_(False)
-        unet.eval()
-
-        self.network = LoRASpecialNetwork(
-            text_encoder=text_encoder,
-            unet=unet,
-            lora_dim=self.network_config.rank,
-            multiplier=1.0,
-            alpha=self.network_config.alpha
-        )
-
-        self.network.force_to(self.device_torch, dtype=dtype)
-
-        self.network.apply_to(
-            text_encoder,
-            unet,
-            self.train_config.train_text_encoder,
-            self.train_config.train_unet
-        )
-
-        self.network.prepare_grad_etc(text_encoder, unet)
-
-        params = self.network.prepare_optimizer_params(
-            text_encoder_lr=self.train_config.lr,
-            unet_lr=self.train_config.lr,
-            default_lr=self.train_config.lr
-        )
-        optimizer_type = self.train_config.optimizer.lower()
-        optimizer = get_optimizer(params, optimizer_type, learning_rate=self.train_config.lr,
-                                  optimizer_params=self.train_config.optimizer_params)
-        lr_scheduler = train_util.get_lr_scheduler(
-            self.train_config.lr_scheduler,
-            optimizer,
-            max_iterations=self.train_config.steps,
-            lr_min=self.train_config.lr / 100,  # not sure why leco did this, but ill do it to
-        )
-        loss_function = torch.nn.MSELoss()
 
+    def hook_before_train_loop(self):
         cache = PromptEmbedsCache()
         prompt_pairs: list[EncodedPromptPair] = []
 
@@ -414,7 +87,7 @@ class TrainSliderProcess(BaseTrainProcess):
                     ]:
                         if cache[prompt] == None:
                             cache[prompt] = train_util.encode_prompts(
-                                tokenizer, text_encoder, [prompt]
+                                self.sd.tokenizer, self.sd.text_encoder, [prompt]
                             )
 
                     # for slider we need to have an enhancer, an eraser, and then
@@ -474,235 +147,184 @@ class TrainSliderProcess(BaseTrainProcess):
                     ]
 
         # move to cpu to save vram
-        # tokenizer.to("cpu")
+        # We don't need text encoder anymore, but keep it on cpu for sampling
-        text_encoder.to("cpu")
+        self.sd.text_encoder.to("cpu")
+        self.prompt_cache = cache
+        self.prompt_pairs = prompt_pairs
         flush()
+        # end hook_before_train_loop
 
-        # sample first
+    def hook_train_loop(self):
-        self.print("Generating baseline samples before training")
+        dtype = get_torch_dtype(self.train_config.dtype)
-        self.sample(0)
+        # get a random pair
+        prompt_pair: EncodedPromptPair = self.prompt_pairs[
+            torch.randint(0, len(self.prompt_pairs), (1,)).item()
+        ]
 
-        self.progress_bar = tqdm(
+        height = prompt_pair.height
-            total=self.train_config.steps,
+        width = prompt_pair.width
-            desc=self.job.name,
+        target_class = prompt_pair.target_class
-            leave=True
+        neutral = prompt_pair.neutral
-        )
+        negative = prompt_pair.negative
-        self.step_num = 0
+        positive = prompt_pair.positive
-        for step in range(self.train_config.steps):
+        weight = prompt_pair.weight
 
-            # get a random pair
+        unet = self.sd.unet
-            prompt_pair: EncodedPromptPair = prompt_pairs[
+        noise_scheduler = self.sd.noise_scheduler
-                torch.randint(0, len(prompt_pairs), (1,)).item()
+        optimizer = self.optimizer
-            ]
+        lr_scheduler = self.lr_scheduler
+        loss_function = torch.nn.MSELoss()
 
-            height = prompt_pair.height
+        # set network multiplier
-            width = prompt_pair.width
+        self.network.multiplier = prompt_pair.multiplier
-            target_class = prompt_pair.target_class
-            neutral = prompt_pair.neutral
-            negative = prompt_pair.negative
-            positive = prompt_pair.positive
-            weight = prompt_pair.weight
 
-            # set network multiplier
+        with torch.no_grad():
-            self.network.multiplier = prompt_pair.multiplier
+            self.sd.noise_scheduler.set_timesteps(
+                self.train_config.max_denoising_steps, device=self.device_torch
+            )
 
-            with torch.no_grad():
+            self.optimizer.zero_grad()
-                noise_scheduler.set_timesteps(
-                    self.train_config.max_denoising_steps, device=self.device_torch
-                )
 
-                optimizer.zero_grad()
+            # ger a random number of steps
+            timesteps_to = torch.randint(
+                1, self.train_config.max_denoising_steps, (1,)
+            ).item()
 
-                # ger a random number of steps
+            # get noise
-                timesteps_to = torch.randint(
+            noise = self.get_latent_noise(
-                    1, self.train_config.max_denoising_steps, (1,)
+                pixel_height=height,
-                ).item()
+                pixel_width=width,
+            ).to(self.device_torch, dtype=dtype)
 
-                # get noise
+            # get latents
-                noise = torch.randn(
+            latents = noise * self.sd.noise_scheduler.init_noise_sigma
-                    (
+            latents = latents.to(self.device_torch, dtype=dtype)
-                        self.train_config.batch_size,
-                        UNET_IN_CHANNELS,
-                        height // VAE_SCALE_FACTOR,
-                        width // VAE_SCALE_FACTOR,
-                    ),
-                    device="cpu",
-                )
-                noise = apply_noise_offset(noise, self.train_config.noise_offset)
-                latents = noise * noise_scheduler.init_noise_sigma
-                latents = latents.to(self.device_torch, dtype=dtype)
-
-                with self.network:
-                    assert self.network.is_active
-                    # A little denoised one is returned
-                    denoised_latents = train_util.diffusion(
-                        unet,
-                        noise_scheduler,
-                        latents,  # pass simple noise latents
-                        train_util.concat_embeddings(
-                            positive,  # unconditional
-                            target_class,  # target
-                            self.train_config.batch_size,
-                        ),
-                        start_timesteps=0,
-                        total_timesteps=timesteps_to,
-                        guidance_scale=3,
-                    )
-
-                noise_scheduler.set_timesteps(1000)
-
-                current_timestep = noise_scheduler.timesteps[
-                    int(timesteps_to * 1000 / self.train_config.max_denoising_steps)
-                ]
-
-                # with network: 0 weight LoRA is enabled outside "with network:"
-                positive_latents = train_util.predict_noise(  # positive_latents
-                    unet,
-                    noise_scheduler,
-                    current_timestep,
-                    denoised_latents,
-                    train_util.concat_embeddings(
-                        positive,  # unconditional
-                        negative,  # positive
-                        self.train_config.batch_size,
-                    ),
-                    guidance_scale=1,
-                ).to("cpu", dtype=torch.float32)
-                neutral_latents = train_util.predict_noise(
-                    unet,
-                    noise_scheduler,
-                    current_timestep,
-                    denoised_latents,
-                    train_util.concat_embeddings(
-                        positive,  # unconditional
-                        neutral,  # neutral
-                        self.train_config.batch_size,
-                    ),
-                    guidance_scale=1,
-                ).to("cpu", dtype=torch.float32)
-                unconditional_latents = train_util.predict_noise(
-                    unet,
-                    noise_scheduler,
-                    current_timestep,
-                    denoised_latents,
-                    train_util.concat_embeddings(
-                        positive,  # unconditional
-                        positive,  # unconditional
-                        self.train_config.batch_size,
-                    ),
-                    guidance_scale=1,
-                ).to("cpu", dtype=torch.float32)
 
             with self.network:
-                target_latents = train_util.predict_noise(
+                assert self.network.is_active
+                # A little denoised one is returned
+                denoised_latents = train_util.diffusion(
                     unet,
                     noise_scheduler,
-                    current_timestep,
+                    latents,  # pass simple noise latents
-                    denoised_latents,
                     train_util.concat_embeddings(
                         positive,  # unconditional
                         target_class,  # target
                         self.train_config.batch_size,
                     ),
-                    guidance_scale=1,
+                    start_timesteps=0,
-                ).to("cpu", dtype=torch.float32)
+                    total_timesteps=timesteps_to,
-
+                    guidance_scale=3,
-                # if self.logging_config.verbose:
-                #     self.print("target_latents:", target_latents[0, 0, :5, :5])
-
-            positive_latents.requires_grad = False
-            neutral_latents.requires_grad = False
-            unconditional_latents.requires_grad = False
-
-            erase = prompt_pair.action == ACTION_TYPES_SLIDER.ERASE_NEGATIVE
-            guidance_scale = 1.0
-
-            offset = guidance_scale * (positive_latents - unconditional_latents)
-
-            offset_neutral = neutral_latents
-            if erase:
-                offset_neutral -= offset
-            else:
-                # enhance
-                offset_neutral += offset
-
-            loss = loss_function(
-                target_latents,
-                offset_neutral,
-            ) * weight
-
-            loss_float = loss.item()
-            if self.train_config.optimizer.startswith('dadaptation'):
-                learning_rate = (
-                        optimizer.param_groups[0]["d"] *
-                        optimizer.param_groups[0]["lr"]
                 )
-            else:
-                learning_rate = optimizer.param_groups[0]['lr']
 
-            self.progress_bar.set_postfix_str(f"lr: {learning_rate:.1e} loss: {loss.item():.3e}")
+            noise_scheduler.set_timesteps(1000)
 
-            loss.backward()
+            current_timestep = noise_scheduler.timesteps[
-            optimizer.step()
+                int(timesteps_to * 1000 / self.train_config.max_denoising_steps)
-            lr_scheduler.step()
+            ]
 
-            del (
+            # with network: 0 weight LoRA is enabled outside "with network:"
-                positive_latents,
+            positive_latents = train_util.predict_noise(  # positive_latents
-                neutral_latents,
+                unet,
-                unconditional_latents,
+                noise_scheduler,
-                target_latents,
+                current_timestep,
-                latents,
+                denoised_latents,
+                train_util.concat_embeddings(
+                    positive,  # unconditional
+                    negative,  # positive
+                    self.train_config.batch_size,
+                ),
+                guidance_scale=1,
+            ).to("cpu", dtype=torch.float32)
+            neutral_latents = train_util.predict_noise(
+                unet,
+                noise_scheduler,
+                current_timestep,
+                denoised_latents,
+                train_util.concat_embeddings(
+                    positive,  # unconditional
+                    neutral,  # neutral
+                    self.train_config.batch_size,
+                ),
+                guidance_scale=1,
+            ).to("cpu", dtype=torch.float32)
+            unconditional_latents = train_util.predict_noise(
+                unet,
+                noise_scheduler,
+                current_timestep,
+                denoised_latents,
+                train_util.concat_embeddings(
+                    positive,  # unconditional
+                    positive,  # unconditional
+                    self.train_config.batch_size,
+                ),
+                guidance_scale=1,
+            ).to("cpu", dtype=torch.float32)
+
+        with self.network:
+            target_latents = train_util.predict_noise(
+                unet,
+                noise_scheduler,
+                current_timestep,
+                denoised_latents,
+                train_util.concat_embeddings(
+                    positive,  # unconditional
+                    target_class,  # target
+                    self.train_config.batch_size,
+                ),
+                guidance_scale=1,
+            ).to("cpu", dtype=torch.float32)
+
+            # if self.logging_config.verbose:
+            #     self.print("target_latents:", target_latents[0, 0, :5, :5])
+
+        positive_latents.requires_grad = False
+        neutral_latents.requires_grad = False
+        unconditional_latents.requires_grad = False
+
+        erase = prompt_pair.action == ACTION_TYPES_SLIDER.ERASE_NEGATIVE
+        guidance_scale = 1.0
+
+        offset = guidance_scale * (positive_latents - unconditional_latents)
+
+        offset_neutral = neutral_latents
+        if erase:
+            offset_neutral -= offset
+        else:
+            # enhance
+            offset_neutral += offset
+
+        loss = loss_function(
+            target_latents,
+            offset_neutral,
+        ) * weight
+
+        loss_float = loss.item()
+        if self.train_config.optimizer.startswith('dadaptation'):
+            learning_rate = (
+                    optimizer.param_groups[0]["d"] *
+                    optimizer.param_groups[0]["lr"]
             )
-            flush()
+        else:
+            learning_rate = optimizer.param_groups[0]['lr']
 
-            # reset network
+        self.progress_bar.set_postfix_str(f"lr: {learning_rate:.1e} loss: {loss.item():.3e}")
-            self.network.multiplier = 1.0
 
-            # don't do on first step
+        loss.backward()
-            if self.step_num != self.start_step:
+        optimizer.step()
-                # pause progress bar
+        lr_scheduler.step()
-                self.progress_bar.unpause()  # makes it so doesn't track time
-                if self.sample_config.sample_every and self.step_num % self.sample_config.sample_every == 0:
-                    # print above the progress bar
-                    self.sample(self.step_num)
-
-                if self.save_config.save_every and self.step_num % self.save_config.save_every == 0:
-                    # print above the progress bar
-                    self.print(f"Saving at step {self.step_num}")
-                    self.save(self.step_num)
-
-                if self.logging_config.log_every and self.step_num % self.logging_config.log_every == 0:
-                    # log to tensorboard
-                    if self.writer is not None:
-                        # get avg loss
-                        self.writer.add_scalar(f"loss", loss_float, self.step_num)
-                        if self.train_config.optimizer.startswith('dadaptation'):
-                            learning_rate = (
-                                    optimizer.param_groups[0]["d"] *
-                                    optimizer.param_groups[0]["lr"]
-                            )
-                        else:
-                            learning_rate = optimizer.param_groups[0]['lr']
-                        self.writer.add_scalar(f"lr", learning_rate, self.step_num)
-                self.progress_bar.refresh()
-
-            # sets progress bar to match out step
-            self.progress_bar.update(step - self.progress_bar.n)
-            # end of step
-            self.step_num = step
-
-        self.sample(self.step_num + 1)
-        print("")
-        self.save()
 
         del (
-            unet,
+            positive_latents,
-            noise_scheduler,
+            neutral_latents,
-            loss,
+            unconditional_latents,
-            optimizer,
+            target_latents,
-            self.network,
+            latents,
-            tokenizer,
-            text_encoder,
         )
-
         flush()
+
+        # reset network
+        self.network.multiplier = 1.0
+
+        return loss_float
+        # end hook_train_loop

toolkit/config_modules.py

@@ -0,0 +1,79 @@
+from typing import List
+
+
+class SaveConfig:
+    def __init__(self, **kwargs):
+        self.save_every: int = kwargs.get('save_every', 1000)
+        self.dtype: str = kwargs.get('save_dtype', 'float16')
+
+
+class LogingConfig:
+    def __init__(self, **kwargs):
+        self.log_every: int = kwargs.get('log_every', 100)
+        self.verbose: bool = kwargs.get('verbose', False)
+        self.use_wandb: bool = kwargs.get('use_wandb', False)
+
+
+class SampleConfig:
+    def __init__(self, **kwargs):
+        self.sample_every: int = kwargs.get('sample_every', 100)
+        self.width: int = kwargs.get('width', 512)
+        self.height: int = kwargs.get('height', 512)
+        self.prompts: list[str] = kwargs.get('prompts', [])
+        self.neg = kwargs.get('neg', False)
+        self.seed = kwargs.get('seed', 0)
+        self.walk_seed = kwargs.get('walk_seed', False)
+        self.guidance_scale = kwargs.get('guidance_scale', 7)
+        self.sample_steps = kwargs.get('sample_steps', 20)
+        self.network_multiplier = kwargs.get('network_multiplier', 1)
+
+
+class NetworkConfig:
+    def __init__(self, **kwargs):
+        self.type: str = kwargs.get('type', 'lierla')
+        self.rank: int = kwargs.get('rank', 4)
+        self.alpha: float = kwargs.get('alpha', 1.0)
+
+
+class TrainConfig:
+    def __init__(self, **kwargs):
+        self.noise_scheduler = kwargs.get('noise_scheduler', 'ddpm')
+        self.steps: int = kwargs.get('steps', 1000)
+        self.lr = kwargs.get('lr', 1e-6)
+        self.optimizer = kwargs.get('optimizer', 'adamw')
+        self.lr_scheduler = kwargs.get('lr_scheduler', 'constant')
+        self.max_denoising_steps: int = kwargs.get('max_denoising_steps', 50)
+        self.batch_size: int = kwargs.get('batch_size', 1)
+        self.dtype: str = kwargs.get('dtype', 'fp32')
+        self.xformers = kwargs.get('xformers', False)
+        self.train_unet = kwargs.get('train_unet', True)
+        self.train_text_encoder = kwargs.get('train_text_encoder', True)
+        self.noise_offset = kwargs.get('noise_offset', 0.0)
+        self.optimizer_params = kwargs.get('optimizer_params', {})
+
+
+class ModelConfig:
+    def __init__(self, **kwargs):
+        self.name_or_path: str = kwargs.get('name_or_path', None)
+        self.is_v2: bool = kwargs.get('is_v2', False)
+        self.is_v_pred: bool = kwargs.get('is_v_pred', False)
+
+        if self.name_or_path is None:
+            raise ValueError('name_or_path must be specified')
+
+
+class SliderTargetConfig:
+    def __init__(self, **kwargs):
+        self.target_class: str = kwargs.get('target_class', '')
+        self.positive: str = kwargs.get('positive', None)
+        self.negative: str = kwargs.get('negative', None)
+        self.multiplier: float = kwargs.get('multiplier', 1.0)
+        self.weight: float = kwargs.get('weight', 1.0)
+
+
+class SliderConfig:
+    def __init__(self, **kwargs):
+        targets = kwargs.get('targets', [])
+        targets = [SliderTargetConfig(**target) for target in targets]
+        self.targets: List[SliderTargetConfig] = targets
+        self.resolutions: List[List[int]] = kwargs.get('resolutions', [[512, 512]])