Added rescaling, locon, sdxl, all kinds of stuff. sdxl is still weird

toolkit/config_modules.py

@@ -5,6 +5,7 @@ class SaveConfig:
     def __init__(self, **kwargs):
         self.save_every: int = kwargs.get('save_every', 1000)
         self.dtype: str = kwargs.get('save_dtype', 'float16')
+        self.max_step_saves_to_keep: int = kwargs.get('max_step_saves_to_keep', 5)
 
 
 class LogingConfig:
@@ -30,8 +31,16 @@ class SampleConfig:
 
 class NetworkConfig:
     def __init__(self, **kwargs):
-        self.type: str = kwargs.get('type', 'lierla')
-        self.rank: int = kwargs.get('rank', 4)
+        self.type: str = kwargs.get('type', 'lora')
+        rank = kwargs.get('rank', None)
+        linear = kwargs.get('linear', None)
+        if rank is not None:
+            self.rank: int = rank # rank for backward compatibility
+            self.linear: int = rank
+        elif linear is not None:
+            self.rank: int = linear
+            self.linear: int = linear
+        self.conv: int = kwargs.get('conv', None)
         self.alpha: float = kwargs.get('alpha', 1.0)
 
 
@@ -51,6 +60,7 @@ class TrainConfig:
         self.noise_offset = kwargs.get('noise_offset', 0.0)
         self.optimizer_params = kwargs.get('optimizer_params', {})
         self.skip_first_sample = kwargs.get('skip_first_sample', False)
+        self.gradient_checkpointing = kwargs.get('gradient_checkpointing', False)
 
 
 class ModelConfig:

toolkit/layers.py

@@ -0,0 +1,31 @@
+import torch
+import torch.nn as nn
+import numpy as np
+
+
+class ReductionKernel(nn.Module):
+    # Tensorflow
+    def __init__(self, in_channels, kernel_size=2, dtype=torch.float32, device=None):
+        if device is None:
+            device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+        super(ReductionKernel, self).__init__()
+        self.kernel_size = kernel_size
+        self.in_channels = in_channels
+        numpy_kernel = self.build_kernel()
+        self.kernel = torch.from_numpy(numpy_kernel).to(device=device, dtype=dtype)
+
+    def build_kernel(self):
+        # tensorflow kernel is  (height, width, in_channels, out_channels)
+        # pytorch kernel is     (out_channels, in_channels, height, width)
+        kernel_size = self.kernel_size
+        channels = self.in_channels
+        kernel_shape = [channels, channels, kernel_size, kernel_size]
+        kernel = np.zeros(kernel_shape, np.float32)
+
+        kernel_value = 1.0 / (kernel_size * kernel_size)
+        for i in range(0, channels):
+            kernel[i, i, :, :] = kernel_value
+        return kernel
+
+    def forward(self, x):
+        return nn.functional.conv2d(x, self.kernel, stride=self.kernel_size, padding=0, groups=1)

toolkit/metadata.py

@@ -1,5 +1,8 @@
 import json
 from collections import OrderedDict
+
+from safetensors import safe_open
+
 from info import software_meta
 
 
@@ -25,4 +28,10 @@ def parse_metadata_from_safetensors(meta: OrderedDict) -> OrderedDict:
             parsed_meta[key] = json.loads(value)
         except json.decoder.JSONDecodeError:
             parsed_meta[key] = value
-    return meta
+    return parsed_meta
+
+
+def load_metadata_from_safetensors(file_path: str) -> OrderedDict:
+    with safe_open(file_path, framework="pt") as f:
+        metadata = f.metadata()
+    return parse_metadata_from_safetensors(metadata)

toolkit/stable_diffusion_model.py

@@ -1,11 +1,18 @@
-from typing import Union
+from typing import Union, OrderedDict
 import sys
 import os
+
+from safetensors.torch import save_file
+
 from toolkit.paths import REPOS_ROOT
+from toolkit.train_tools import get_torch_dtype
+
 sys.path.append(REPOS_ROOT)
 sys.path.append(os.path.join(REPOS_ROOT, 'leco'))
 from leco import train_util
 import torch
+from library import model_util
+from library.sdxl_model_util import convert_text_encoder_2_state_dict_to_sdxl
 
 
 class PromptEmbeds:
@@ -22,6 +29,12 @@ class PromptEmbeds:
             self.text_embeds = args
             self.pooled_embeds = None
 
+    def to(self, **kwargs):
+        self.text_embeds = self.text_embeds.to(**kwargs)
+        if self.pooled_embeds is not None:
+            self.pooled_embeds = self.pooled_embeds.to(**kwargs)
+        return self
+
 
 class StableDiffusion:
     def __init__(
@@ -61,3 +74,41 @@ class StableDiffusion:
                     self.tokenizer, self.text_encoder, prompt
                 )
             )
+
+    def save(self, output_file: str, meta: OrderedDict, save_dtype=get_torch_dtype('fp16'), logit_scale=None):
+        # todo see what logit scale is
+        if self.is_xl:
+
+            state_dict = {}
+
+            def update_sd(prefix, sd):
+                for k, v in sd.items():
+                    key = prefix + k
+                    v = v.detach().clone().to("cpu").to(get_torch_dtype(save_dtype))
+                    state_dict[key] = v
+
+            # Convert the UNet model
+            update_sd("model.diffusion_model.", self.unet.state_dict())
+
+            # Convert the text encoders
+            update_sd("conditioner.embedders.0.transformer.", self.text_encoder[0].state_dict())
+
+            text_enc2_dict = convert_text_encoder_2_state_dict_to_sdxl(self.text_encoder[1].state_dict(), logit_scale)
+            update_sd("conditioner.embedders.1.model.", text_enc2_dict)
+
+            # Convert the VAE
+            vae_dict = model_util.convert_vae_state_dict(self.vae.state_dict())
+            update_sd("first_stage_model.", vae_dict)
+
+            # Put together new checkpoint
+            key_count = len(state_dict.keys())
+            new_ckpt = {"state_dict": state_dict}
+
+            if model_util.is_safetensors(output_file):
+                save_file(state_dict, output_file)
+            else:
+                torch.save(new_ckpt, output_file, meta)
+
+            return key_count
+        else:
+            raise NotImplementedError("sdv1.x, sdv2.x is not implemented yet")

toolkit/train_tools.py

@@ -2,6 +2,7 @@ import argparse
 import json
 import os
 import time
+from typing import TYPE_CHECKING
 
 from diffusers import (
     StableDiffusionPipeline,
@@ -21,8 +22,6 @@ from library.lpw_stable_diffusion import StableDiffusionLongPromptWeightingPipel
 import torch
 import re
 
-from toolkit.stable_diffusion_model import PromptEmbeds
-
 SCHEDULER_LINEAR_START = 0.00085
 SCHEDULER_LINEAR_END = 0.0120
 SCHEDULER_TIMESTEPS = 1000
@@ -381,11 +380,16 @@ def apply_noise_offset(noise, noise_offset):
     return noise
 
 
+if TYPE_CHECKING:
+    from toolkit.stable_diffusion_model import PromptEmbeds
+
+
 def concat_prompt_embeddings(
-        unconditional: PromptEmbeds,
-        conditional: PromptEmbeds,
+        unconditional: 'PromptEmbeds',
+        conditional: 'PromptEmbeds',
         n_imgs: int,
 ):
+    from toolkit.stable_diffusion_model import PromptEmbeds
     text_embeds = torch.cat(
         [unconditional.text_embeds, conditional.text_embeds]
     ).repeat_interleave(n_imgs, dim=0)