Added support for training ssd-1B. Added support for saving models into diffusers format. We can currently save in safetensors format for ssd-1b, but diffusers cannot load it yet.

toolkit/config_modules.py

@@ -9,12 +9,17 @@ from toolkit.prompt_utils import PromptEmbeds
 
 ImgExt = Literal['jpg', 'png', 'webp']
 
+SaveFormat = Literal['safetensors', 'diffusers']
+
 
 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)
+        self.save_format: SaveFormat = kwargs.get('save_format', 'safetensors')
+        if self.save_format not in ['safetensors', 'diffusers']:
+            raise ValueError(f"save_format must be safetensors or diffusers, got {self.save_format}")
 
 
 class LogingConfig:
@@ -187,7 +192,7 @@ class TrainConfig:
         self.img_multiplier = kwargs.get('img_multiplier', 1.0)
 
         # match the norm of the noise before computing loss. This will help the model maintain its
-        #current understandin of the brightness of images.
+        # current understandin of the brightness of images.
 
         self.match_noise_norm = kwargs.get('match_noise_norm', False)
 
@@ -229,6 +234,7 @@ class ModelConfig:
         self.name_or_path: str = kwargs.get('name_or_path', None)
         self.is_v2: bool = kwargs.get('is_v2', False)
         self.is_xl: bool = kwargs.get('is_xl', False)
+        self.is_ssd: bool = kwargs.get('is_ssd', False)
         self.is_v_pred: bool = kwargs.get('is_v_pred', False)
         self.dtype: str = kwargs.get('dtype', 'float16')
         self.vae_path = kwargs.get('vae_path', None)
@@ -242,6 +248,10 @@ class ModelConfig:
         if self.name_or_path is None:
             raise ValueError('name_or_path must be specified')
 
+        if self.is_ssd:
+            # sed sdxl as true since it is mostly the same architecture
+            self.is_xl = True
+
 
 class ReferenceDatasetConfig:
     def __init__(self, **kwargs):

toolkit/keymaps/stable_diffusion_ssd_unmatched.json

@@ -0,0 +1,21 @@
+{
+    "ldm": {
+        "conditioner.embedders.0.transformer.text_model.embeddings.position_ids": {
+            "shape": [
+                1,
+                77
+            ],
+            "min": 0.0,
+            "max": 76.0
+        },
+        "conditioner.embedders.1.model.text_model.embeddings.position_ids": {
+            "shape": [
+                1,
+                77
+            ],
+            "min": 0.0,
+            "max": 76.0
+        }
+    },
+    "diffusers": {}
+}

toolkit/saving.py

@@ -97,7 +97,7 @@ def convert_state_dict_to_ldm_with_mapping(
 
 def get_ldm_state_dict_from_diffusers(
         state_dict: 'OrderedDict',
-        sd_version: Literal['1', '2', 'sdxl'] = '2',
+        sd_version: Literal['1', '2', 'sdxl', 'ssd'] = '2',
         device='cpu',
         dtype=get_torch_dtype('fp32'),
 ):
@@ -111,6 +111,10 @@ def get_ldm_state_dict_from_diffusers(
         # load our base
         base_path = os.path.join(KEYMAPS_ROOT, 'stable_diffusion_sdxl_ldm_base.safetensors')
         mapping_path = os.path.join(KEYMAPS_ROOT, 'stable_diffusion_sdxl.json')
+    elif sd_version == 'ssd':
+        # load our base
+        base_path = os.path.join(KEYMAPS_ROOT, 'stable_diffusion_ssd_ldm_base.safetensors')
+        mapping_path = os.path.join(KEYMAPS_ROOT, 'stable_diffusion_ssd.json')
     else:
         raise ValueError(f"Invalid sd_version {sd_version}")
 
@@ -129,7 +133,7 @@ def save_ldm_model_from_diffusers(
         output_file: str,
         meta: 'OrderedDict',
         save_dtype=get_torch_dtype('fp16'),
-        sd_version: Literal['1', '2', 'sdxl'] = '2'
+        sd_version: Literal['1', '2', 'sdxl', 'ssd'] = '2'
 ):
     converted_state_dict = get_ldm_state_dict_from_diffusers(
         sd.state_dict(),
@@ -148,13 +152,14 @@ def save_lora_from_diffusers(
         output_file: str,
         meta: 'OrderedDict',
         save_dtype=get_torch_dtype('fp16'),
-        sd_version: Literal['1', '2', 'sdxl'] = '2'
+        sd_version: Literal['1', '2', 'sdxl', 'ssd'] = '2'
 ):
     converted_state_dict = OrderedDict()
     # only handle sxdxl for now
-    if sd_version != 'sdxl':
+    if sd_version != 'sdxl' and sd_version != 'ssd':
         raise ValueError(f"Invalid sd_version {sd_version}")
     for key, value in lora_state_dict.items():
+        # todo verify if this works with ssd
         # test encoders share keys for some reason
         if key.begins_with('lora_te'):
             converted_state_dict[key] = value.detach().to('cpu', dtype=save_dtype)

toolkit/stable_diffusion_model.py

@@ -9,6 +9,7 @@ import sys
 import os
 from collections import OrderedDict
 
+import yaml
 from PIL import Image
 from diffusers.pipelines.stable_diffusion_xl.pipeline_stable_diffusion_xl import rescale_noise_cfg
 from safetensors.torch import save_file, load_file
@@ -125,6 +126,7 @@ class StableDiffusion:
         self.adapter: Union['T2IAdapter', 'IPAdapter', None] = None
         self.is_xl = model_config.is_xl
         self.is_v2 = model_config.is_v2
+        self.is_ssd = model_config.is_ssd
 
         self.use_text_encoder_1 = model_config.use_text_encoder_1
         self.use_text_encoder_2 = model_config.use_text_encoder_2
@@ -157,7 +159,7 @@ class StableDiffusion:
         if self.model_config.vae_path is not None:
             load_args['vae'] = load_vae(self.model_config.vae_path, dtype)
 
-        if self.model_config.is_xl:
+        if self.model_config.is_xl or self.model_config.is_ssd:
             if self.custom_pipeline is not None:
                 pipln = self.custom_pipeline
             else:
@@ -165,7 +167,7 @@ class StableDiffusion:
                 # pipln = StableDiffusionKDiffusionXLPipeline
 
             # see if path exists
-            if not os.path.exists(model_path):
+            if not os.path.exists(model_path) or os.path.isdir(model_path):
                 # try to load with default diffusers
                 pipe = pipln.from_pretrained(
                     model_path,
@@ -176,19 +178,11 @@ class StableDiffusion:
                     **load_args
                 )
             else:
-                try:
-                    pipe = pipln.from_single_file(
-                        model_path,
-                        device=self.device_torch,
-                        torch_dtype=self.torch_dtype,
-                    )
-                except Exception as e:
-                    print("Error loading model from single file. Trying to load from pretrained")
-                    pipe = pipln.from_pretrained(
-                        model_path,
-                        device=self.device_torch,
-                        torch_dtype=self.torch_dtype,
-                    )
+                pipe = pipln.from_single_file(
+                    model_path,
+                    device=self.device_torch,
+                    torch_dtype=self.torch_dtype,
+                )
             flush()
 
             text_encoders = [pipe.text_encoder, pipe.text_encoder_2]
@@ -900,17 +894,34 @@ class StableDiffusion:
             version_string = '2'
         if self.is_xl:
             version_string = 'sdxl'
-        save_ldm_model_from_diffusers(
-            sd=self,
-            output_file=output_file,
-            meta=meta,
-            save_dtype=save_dtype,
-            sd_version=version_string,
-        )
-        if self.config_file is not None:
-            output_path_no_ext = os.path.splitext(output_file)[0]
-            output_config_path = f"{output_path_no_ext}.yaml"
-            shutil.copyfile(self.config_file, output_config_path)
+        if self.is_ssd:
+            # overwrite sdxl because both wil be true here
+            version_string = 'ssd'
+        # if output file does not end in .safetensors, then it is a directory and we are
+        # saving in diffusers format
+        if not output_file.endswith('.safetensors'):
+            # diffusers
+            self.pipeline.save_pretrained(
+                save_directory=output_file,
+                safe_serialization=True,
+            )
+            # save out meta config
+            meta_path = os.path.join(output_file, 'aitk_meta.yaml')
+            with open(meta_path, 'w') as f:
+                yaml.dump(meta, f)
+
+        else:
+            save_ldm_model_from_diffusers(
+                sd=self,
+                output_file=output_file,
+                meta=meta,
+                save_dtype=save_dtype,
+                sd_version=version_string,
+            )
+            if self.config_file is not None:
+                output_path_no_ext = os.path.splitext(output_file)[0]
+                output_config_path = f"{output_path_no_ext}.yaml"
+                shutil.copyfile(self.config_file, output_config_path)
 
     def prepare_optimizer_params(
             self,