Merge branch 'main' into development

toolkit/keymaps/stable_diffusion_sd2_unmatched.json

@@ -0,0 +1,200 @@
+{
+    "ldm": {
+        "alphas_cumprod": {
+            "shape": [
+                1000
+            ],
+            "min": 0.00466156005859375,
+            "max": 0.9990234375
+        },
+        "alphas_cumprod_prev": {
+            "shape": [
+                1000
+            ],
+            "min": 0.0047149658203125,
+            "max": 1.0
+        },
+        "betas": {
+            "shape": [
+                1000
+            ],
+            "min": 0.0008502006530761719,
+            "max": 0.01200103759765625
+        },
+        "cond_stage_model.model.logit_scale": {
+            "shape": [],
+            "min": 4.60546875,
+            "max": 4.60546875
+        },
+        "cond_stage_model.model.text_projection": {
+            "shape": [
+                1024,
+                1024
+            ],
+            "min": -0.109130859375,
+            "max": 0.09271240234375
+        },
+        "cond_stage_model.model.transformer.resblocks.23.attn.in_proj_bias": {
+            "shape": [
+                3072
+            ],
+            "min": -2.525390625,
+            "max": 2.591796875
+        },
+        "cond_stage_model.model.transformer.resblocks.23.attn.in_proj_weight": {
+            "shape": [
+                3072,
+                1024
+            ],
+            "min": -0.12261962890625,
+            "max": 0.1258544921875
+        },
+        "cond_stage_model.model.transformer.resblocks.23.attn.out_proj.bias": {
+            "shape": [
+                1024
+            ],
+            "min": -0.422607421875,
+            "max": 1.17578125
+        },
+        "cond_stage_model.model.transformer.resblocks.23.attn.out_proj.weight": {
+            "shape": [
+                1024,
+                1024
+            ],
+            "min": -0.0738525390625,
+            "max": 0.08673095703125
+        },
+        "cond_stage_model.model.transformer.resblocks.23.ln_1.bias": {
+            "shape": [
+                1024
+            ],
+            "min": -3.392578125,
+            "max": 0.90625
+        },
+        "cond_stage_model.model.transformer.resblocks.23.ln_1.weight": {
+            "shape": [
+                1024
+            ],
+            "min": 0.379638671875,
+            "max": 2.02734375
+        },
+        "cond_stage_model.model.transformer.resblocks.23.ln_2.bias": {
+            "shape": [
+                1024
+            ],
+            "min": -0.833984375,
+            "max": 2.525390625
+        },
+        "cond_stage_model.model.transformer.resblocks.23.ln_2.weight": {
+            "shape": [
+                1024
+            ],
+            "min": 1.17578125,
+            "max": 2.037109375
+        },
+        "cond_stage_model.model.transformer.resblocks.23.mlp.c_fc.bias": {
+            "shape": [
+                4096
+            ],
+            "min": -1.619140625,
+            "max": 0.5595703125
+        },
+        "cond_stage_model.model.transformer.resblocks.23.mlp.c_fc.weight": {
+            "shape": [
+                4096,
+                1024
+            ],
+            "min": -0.08953857421875,
+            "max": 0.13232421875
+        },
+        "cond_stage_model.model.transformer.resblocks.23.mlp.c_proj.bias": {
+            "shape": [
+                1024
+            ],
+            "min": -1.8662109375,
+            "max": 0.74658203125
+        },
+        "cond_stage_model.model.transformer.resblocks.23.mlp.c_proj.weight": {
+            "shape": [
+                1024,
+                4096
+            ],
+            "min": -0.12939453125,
+            "max": 0.1009521484375
+        },
+        "log_one_minus_alphas_cumprod": {
+            "shape": [
+                1000
+            ],
+            "min": -7.0703125,
+            "max": -0.004669189453125
+        },
+        "model_ema.decay": {
+            "shape": [],
+            "min": 1.0,
+            "max": 1.0
+        },
+        "model_ema.num_updates": {
+            "shape": [],
+            "min": 219996,
+            "max": 219996
+        },
+        "posterior_log_variance_clipped": {
+            "shape": [
+                1000
+            ],
+            "min": -46.0625,
+            "max": -4.421875
+        },
+        "posterior_mean_coef1": {
+            "shape": [
+                1000
+            ],
+            "min": 0.000827789306640625,
+            "max": 1.0
+        },
+        "posterior_mean_coef2": {
+            "shape": [
+                1000
+            ],
+            "min": 0.0,
+            "max": 0.99560546875
+        },
+        "posterior_variance": {
+            "shape": [
+                1000
+            ],
+            "min": 0.0,
+            "max": 0.01200103759765625
+        },
+        "sqrt_alphas_cumprod": {
+            "shape": [
+                1000
+            ],
+            "min": 0.0682373046875,
+            "max": 0.99951171875
+        },
+        "sqrt_one_minus_alphas_cumprod": {
+            "shape": [
+                1000
+            ],
+            "min": 0.0291595458984375,
+            "max": 0.99755859375
+        },
+        "sqrt_recip_alphas_cumprod": {
+            "shape": [
+                1000
+            ],
+            "min": 1.0,
+            "max": 14.6484375
+        },
+        "sqrt_recipm1_alphas_cumprod": {
+            "shape": [
+                1000
+            ],
+            "min": 0.0291595458984375,
+            "max": 14.6171875
+        }
+    },
+    "diffusers": {}
+}

toolkit/saving.py

@@ -56,7 +56,7 @@ def convert_state_dict_to_ldm_with_mapping(
             tensor_to_slice = diffusers_state_dict[ldm_diffusers_operator_map[ldm_key]['slice'][0]]
             slice_text = diffusers_state_dict[ldm_diffusers_operator_map[ldm_key]['slice'][1]]
             converted_state_dict[ldm_key] = tensor_to_slice[get_slices_from_string(slice_text)].detach().to(device,
-                                                                                                             dtype=dtype)
+                                                                                                            dtype=dtype)
 
     # process the rest of the keys
     for ldm_key in ldm_diffusers_keymap:
@@ -71,6 +71,35 @@ def convert_state_dict_to_ldm_with_mapping(
     return converted_state_dict
 
 
+def get_ldm_state_dict_from_diffusers(
+        state_dict: 'OrderedDict',
+        sd_version: Literal['1', '2', 'sdxl'] = '2',
+        device='cpu',
+        dtype=get_torch_dtype('fp32'),
+):
+    if sd_version == '1':
+        base_path = os.path.join(KEYMAPS_ROOT, 'stable_diffusion_sd1_ldm_base.safetensors')
+        mapping_path = os.path.join(KEYMAPS_ROOT, 'stable_diffusion_sd1.json')
+    elif sd_version == '2':
+        base_path = os.path.join(KEYMAPS_ROOT, 'stable_diffusion_sd2_ldm_base.safetensors')
+        mapping_path = os.path.join(KEYMAPS_ROOT, 'stable_diffusion_sd2.json')
+    elif sd_version == 'sdxl':
+        # 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')
+    else:
+        raise ValueError(f"Invalid sd_version {sd_version}")
+
+        # convert the state dict
+    return convert_state_dict_to_ldm_with_mapping(
+        state_dict,
+        mapping_path,
+        base_path,
+        device=device,
+        dtype=dtype
+    )
+
+
 def save_ldm_model_from_diffusers(
         sd: 'StableDiffusion',
         output_file: str,
@@ -78,21 +107,13 @@ def save_ldm_model_from_diffusers(
         save_dtype=get_torch_dtype('fp16'),
         sd_version: Literal['1', '2', 'sdxl'] = '2'
 ):
-    if sd_version != 'sdxl':
-        # not supported yet
-        raise NotImplementedError("Only SDXL is supported at this time with this method")
-    # 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')
-
-    # convert the state dict
-    converted_state_dict = convert_state_dict_to_ldm_with_mapping(
+    converted_state_dict = get_ldm_state_dict_from_diffusers(
         sd.state_dict(),
-        mapping_path,
-        base_path,
+        sd_version,
         device='cpu',
         dtype=save_dtype
     )
+
     # make sure parent folder exists
     os.makedirs(os.path.dirname(output_file), exist_ok=True)
     save_file(converted_state_dict, output_file, metadata=meta)

toolkit/stable_diffusion_model.py

@@ -709,61 +709,28 @@ class StableDiffusion:
                 for i, encoder in enumerate(self.text_encoder):
                     for k, v in encoder.state_dict().items():
                         new_key = k if k.startswith(
-                            f"{SD_PREFIX_TEXT_ENCODER}{i}") else f"{SD_PREFIX_TEXT_ENCODER}{i}_{k}"
+                            f"{SD_PREFIX_TEXT_ENCODER}{i}_") else f"{SD_PREFIX_TEXT_ENCODER}{i}_{k}"
                         state_dict[new_key] = v
             else:
                 for k, v in self.text_encoder.state_dict().items():
-                    new_key = k if k.startswith(f"{SD_PREFIX_TEXT_ENCODER}") else f"{SD_PREFIX_TEXT_ENCODER}_{k}"
+                    new_key = k if k.startswith(f"{SD_PREFIX_TEXT_ENCODER}_") else f"{SD_PREFIX_TEXT_ENCODER}_{k}"
                     state_dict[new_key] = v
         if unet:
             for k, v in self.unet.state_dict().items():
-                new_key = k if k.startswith(f"{SD_PREFIX_UNET}") else f"{SD_PREFIX_UNET}_{k}"
+                new_key = k if k.startswith(f"{SD_PREFIX_UNET}_") else f"{SD_PREFIX_UNET}_{k}"
                 state_dict[new_key] = v
         return state_dict
 
     def save(self, output_file: str, meta: OrderedDict, save_dtype=get_torch_dtype('fp16'), logit_scale=None):
-        state_dict = {}
-        # prepare metadata
-        meta = get_meta_for_safetensors(meta)
-
-        def update_sd(prefix, sd):
-            for k, v in sd.items():
-                key = prefix + k
-                v = v.detach().clone()
-                state_dict[key] = v.to("cpu", dtype=get_torch_dtype(save_dtype))
-                # make sure there are not nan values
-                if torch.isnan(state_dict[key]).any():
-                    raise ValueError(f"NaN value in state dict: {key}")
-
-        # todo see what logit scale is
+        version_string = '1'
+        if self.is_v2:
+            version_string = '2'
         if self.is_xl:
-            save_ldm_model_from_diffusers(
-                sd=self,
-                output_file=output_file,
-                meta=meta,
-                save_dtype=save_dtype,
-                sd_version='sdxl',
-            )
-
-        else:
-            # Convert the UNet model
-            unet_state_dict = convert_unet_state_dict_to_sd(self.is_v2, self.unet.state_dict())
-            update_sd("model.diffusion_model.", unet_state_dict)
-
-            # Convert the text encoder model
-            if self.is_v2:
-                make_dummy = True
-                text_enc_dict = convert_text_encoder_state_dict_to_sd_v2(self.text_encoder.state_dict(), make_dummy)
-                update_sd("cond_stage_model.model.", text_enc_dict)
-            else:
-                text_enc_dict = self.text_encoder.state_dict()
-                update_sd("cond_stage_model.transformer.", text_enc_dict)
-
-                # Convert the VAE
-            if self.vae is not None:
-                vae_dict = model_util.convert_vae_state_dict(self.vae.state_dict())
-                update_sd("first_stage_model.", vae_dict)
-
-            # make sure parent folder exists
-            os.makedirs(os.path.dirname(output_file), exist_ok=True)
-            save_file(state_dict, output_file, metadata=meta)
+            version_string = 'sdxl'
+        save_ldm_model_from_diffusers(
+            sd=self,
+            output_file=output_file,
+            meta=meta,
+            save_dtype=save_dtype,
+            sd_version=version_string,
+        )