



'''simple docstring'''





import gzip
import hashlib
import json
import multiprocessing
import os
import re
import shutil
import time
from pathlib import Path

import numpy as np
from arguments import PreprocessingArguments
from datasets import load_dataset
from minhash_deduplication import deduplicate_dataset

from transformers import AutoTokenizer, HfArgumentParser


_lowerCAmelCase      						=	re.compile(R'''\s+''')
def    _SCREAMING_SNAKE_CASE  (       UpperCamelCase			):



    """simple docstring"""





    return {"hash": hashlib.mda(re.sub(UpperCamelCase							,    """"""							,    example["""content"""]			).encode("""utf-8"""			)			).hexdigest()}
def    _SCREAMING_SNAKE_CASE  (       UpperCamelCase			):



    """simple docstring"""





    lowerCAmelCase__ :  str							      =					[len(UpperCamelCase			) for line in example["""content"""].splitlines()]
    return {"line_mean": np.mean(UpperCamelCase			), "line_max": max(UpperCamelCase			)}
def    _SCREAMING_SNAKE_CASE  (       UpperCamelCase			):



    """simple docstring"""





    lowerCAmelCase__ :  Any							      =					np.mean([c.isalnum() for c in example["""content"""]]			)
    return {"alpha_frac": alpha_frac}
def    _SCREAMING_SNAKE_CASE  (       UpperCamelCase							,    UpperCamelCase			):



    """simple docstring"""





    if example["hash"] in uniques:
        uniques.remove(example["""hash"""]			)
        return True
    else:
        return False
def    _SCREAMING_SNAKE_CASE  (       UpperCamelCase							,    UpperCamelCase=5			):



    """simple docstring"""





    lowerCAmelCase__ :  Union[str, Any]							      =					["""auto-generated""", """autogenerated""", """automatically generated"""]
    lowerCAmelCase__ :  Any							      =					example["""content"""].splitlines()
    for _, line in zip(range(UpperCamelCase			)							,    UpperCamelCase			):
        for keyword in keywords:
            if keyword in line.lower():
                return {"autogenerated": True}
    else:
        return {"autogenerated": False}
def    _SCREAMING_SNAKE_CASE  (       UpperCamelCase							,    UpperCamelCase=5							,    UpperCamelCase=0.05			):



    """simple docstring"""





    lowerCAmelCase__ :  Dict							      =					["""unit tests""", """test file""", """configuration file"""]
    lowerCAmelCase__ :  str							      =					example["""content"""].splitlines()
    lowerCAmelCase__ :  str							      =					0
    lowerCAmelCase__ :  Union[str, Any]							      =					0
    # first test
    for _, line in zip(range(UpperCamelCase			)							,    UpperCamelCase			):
        for keyword in keywords:
            if keyword in line.lower():
                return {"config_or_test": True}
    # second test
    lowerCAmelCase__ :  Any							      =					example["""content"""].count("""\n"""			)
    lowerCAmelCase__ :  Optional[Any]							      =					int(coeff * nlines			)
    for line in lines:
        count_config += line.lower().count("""config"""			)
        count_test += line.lower().count("""test"""			)
        if count_config > threshold or count_test > threshold:
            return {"config_or_test": True}
    return {"config_or_test": False}
def    _SCREAMING_SNAKE_CASE  (       UpperCamelCase			):



    """simple docstring"""





    lowerCAmelCase__ :  str							      =					["""def """, """class """, """for """, """while """]
    lowerCAmelCase__ :  Union[str, Any]							      =					example["""content"""].splitlines()
    for line in lines:
        for keyword in keywords:
            if keyword in line.lower():
                return {"has_no_keywords": False}
    return {"has_no_keywords": True}
def    _SCREAMING_SNAKE_CASE  (       UpperCamelCase							,    UpperCamelCase=4			):



    """simple docstring"""





    lowerCAmelCase__ :  Union[str, Any]							      =					example["""content"""].splitlines()
    lowerCAmelCase__ :  List[Any]							      =					0
    for line in lines:
        counter += line.lower().count("""="""			)
        if counter > minimum:
            return {"has_few_assignments": False}
    return {"has_few_assignments": True}
def    _SCREAMING_SNAKE_CASE  (       UpperCamelCase			):



    """simple docstring"""





    lowerCAmelCase__ :  List[str]							      =					tokenizer(example["""content"""]							,    truncation=UpperCamelCase			)["""input_ids"""]
    lowerCAmelCase__ :  Union[str, Any]							      =					len(example["""content"""]			) / len(UpperCamelCase			)
    return {"ratio": ratio}
def    _SCREAMING_SNAKE_CASE  (       UpperCamelCase			):



    """simple docstring"""





    lowerCAmelCase__ :  int							      =					{}
    results.update(get_hash(UpperCamelCase			)			)
    results.update(line_stats(UpperCamelCase			)			)
    results.update(alpha_stats(UpperCamelCase			)			)
    results.update(char_token_ratio(UpperCamelCase			)			)
    results.update(is_autogenerated(UpperCamelCase			)			)
    results.update(is_config_or_test(UpperCamelCase			)			)
    results.update(has_no_keywords(UpperCamelCase			)			)
    results.update(has_few_assignments(UpperCamelCase			)			)
    return results
def    _SCREAMING_SNAKE_CASE  (       UpperCamelCase							,    UpperCamelCase							,    UpperCamelCase			):



    """simple docstring"""





    if not check_uniques(UpperCamelCase							,    UpperCamelCase			):
        return False
    elif example["autogenerated"]:
        return False
    elif example["line_max"] > args.line_max:
        return False
    elif example["line_mean"] > args.line_mean:
        return False
    elif example["alpha_frac"] < args.alpha_frac:
        return False
    elif example["ratio"] < args.min_token_ratio:
        return False
    elif example["config_or_test"] and np.random.rand() <= args.filter_proba:
        return False
    elif example["has_no_keywords"] and np.random.rand() <= args.filter_proba:
        return False
    elif example["has_few_assignments"]:
        return False
    else:
        return True
def    _SCREAMING_SNAKE_CASE  (       UpperCamelCase			):



    """simple docstring"""





    with open(UpperCamelCase							,    """rb"""			) as f_in:
        with gzip.open(str(UpperCamelCase			) + """.gz"""							,    """wb"""							,    compresslevel=6			) as f_out:
            shutil.copyfileobj(UpperCamelCase							,    UpperCamelCase			)
    os.unlink(UpperCamelCase			)


# Settings
_lowerCAmelCase      						=	HfArgumentParser(PreprocessingArguments)
_lowerCAmelCase      						=	parser.parse_args()
if args.num_workers is None:
 _lowerCAmelCase      						=	multiprocessing.cpu_count()
_lowerCAmelCase      						=	AutoTokenizer.from_pretrained(args.tokenizer_dir)

# Load dataset
_lowerCAmelCase      						=	time.time()
_lowerCAmelCase      						=	load_dataset(args.dataset_name, split='''train''')
print(F"""Time to load dataset: {time.time()-t_start:.2f}""")

# Run preprocessing
_lowerCAmelCase      						=	time.time()
_lowerCAmelCase      						=	ds.map(preprocess, num_proc=args.num_workers)
print(F"""Time to preprocess dataset: {time.time()-t_start:.2f}""")

# Deduplicate hashes
_lowerCAmelCase      						=	set(ds.unique('''hash'''))
_lowerCAmelCase      						=	len(uniques) / len(ds)
print(F"""Fraction of duplicates: {1-frac:.2%}""")

# Deduplicate data and apply heuristics
_lowerCAmelCase      						=	time.time()
_lowerCAmelCase      						=	ds.filter(filter, fn_kwargs={'''uniques''': uniques, '''args''': args})
print(F"""Time to filter dataset: {time.time()-t_start:.2f}""")
print(F"""Size of filtered dataset: {len(ds_filter)}""")

# Deduplicate with minhash and jaccard similarity
if args.near_deduplication:
 _lowerCAmelCase      						=	time.time()
 _lowerCAmelCase						,			_lowerCAmelCase      						=	deduplicate_dataset(ds_filter, args.jaccard_threshold)
 print(F"""Time to deduplicate dataset: {time.time()-t_start:.2f}""")
 print(F"""Size of deduplicate dataset: {len(ds_filter)}""")

# Save data in batches of samples_per_file
_lowerCAmelCase      						=	Path(args.output_dir)
output_dir.mkdir(exist_ok=True)

# save duplicate_clusters in the output_dir as artifacts
# not sure it is the right place the save it
if args.near_deduplication:
 with open(output_dir / '''duplicate_clusters.json''', '''w''') as f:
  json.dump(duplicate_clusters, f)

_lowerCAmelCase      						=	output_dir / '''data'''
data_dir.mkdir(exist_ok=True)

_lowerCAmelCase      						=	time.time()
for file_number, index in enumerate(range(0, len(ds_filter), args.samples_per_file)):
 _lowerCAmelCase      						=	str(data_dir / F"""file-{file_number+1:012}.json""")
 _lowerCAmelCase      						=	min(len(ds_filter), index + args.samples_per_file)
 ds_filter.select(list(range(index, end_index))).to_json(file_path)
 compress_file(file_path)




print(F"""Time to save dataset: {time.time()-t_start:.2f}""")







'''simple docstring'''





import argparse
import os.path as osp
import re

import torch
from safetensors.torch import load_file, save_file


# =================#
# UNet Conversion #
# =================#

_lowerCAmelCase      						=	[
    # (stable-diffusion, HF Diffusers)
    ('''time_embed.0.weight''', '''time_embedding.linear_1.weight'''),
    ('''time_embed.0.bias''', '''time_embedding.linear_1.bias'''),
    ('''time_embed.2.weight''', '''time_embedding.linear_2.weight'''),
    ('''time_embed.2.bias''', '''time_embedding.linear_2.bias'''),
    ('''input_blocks.0.0.weight''', '''conv_in.weight'''),
    ('''input_blocks.0.0.bias''', '''conv_in.bias'''),
    ('''out.0.weight''', '''conv_norm_out.weight'''),
    ('''out.0.bias''', '''conv_norm_out.bias'''),
    ('''out.2.weight''', '''conv_out.weight'''),
    ('''out.2.bias''', '''conv_out.bias'''),
]

_lowerCAmelCase      						=	[
    # (stable-diffusion, HF Diffusers)
    ('''in_layers.0''', '''norm1'''),
    ('''in_layers.2''', '''conv1'''),
    ('''out_layers.0''', '''norm2'''),
    ('''out_layers.3''', '''conv2'''),
    ('''emb_layers.1''', '''time_emb_proj'''),
    ('''skip_connection''', '''conv_shortcut'''),
]

_lowerCAmelCase      						=	[]
# hardcoded number of downblocks and resnets/attentions...
# would need smarter logic for other networks.
for i in range(4):
 # loop over downblocks/upblocks

 for j in range(2):
  # loop over resnets/attentions for downblocks
  _lowerCAmelCase      						=	F"""down_blocks.{i}.resnets.{j}."""
  _lowerCAmelCase      						=	F"""input_blocks.{3*i + j + 1}.0."""
  unet_conversion_map_layer.append((sd_down_res_prefix, hf_down_res_prefix))

  if i < 3:
   # no attention layers in down_blocks.3
   _lowerCAmelCase      						=	F"""down_blocks.{i}.attentions.{j}."""
   _lowerCAmelCase      						=	F"""input_blocks.{3*i + j + 1}.1."""
   unet_conversion_map_layer.append((sd_down_atn_prefix, hf_down_atn_prefix))

 for j in range(3):
  # loop over resnets/attentions for upblocks
  _lowerCAmelCase      						=	F"""up_blocks.{i}.resnets.{j}."""
  _lowerCAmelCase      						=	F"""output_blocks.{3*i + j}.0."""
  unet_conversion_map_layer.append((sd_up_res_prefix, hf_up_res_prefix))

  if i > 0:
   # no attention layers in up_blocks.0
   _lowerCAmelCase      						=	F"""up_blocks.{i}.attentions.{j}."""
   _lowerCAmelCase      						=	F"""output_blocks.{3*i + j}.1."""
   unet_conversion_map_layer.append((sd_up_atn_prefix, hf_up_atn_prefix))

 if i < 3:
  # no downsample in down_blocks.3
  _lowerCAmelCase      						=	F"""down_blocks.{i}.downsamplers.0.conv."""
  _lowerCAmelCase      						=	F"""input_blocks.{3*(i+1)}.0.op."""
  unet_conversion_map_layer.append((sd_downsample_prefix, hf_downsample_prefix))

  # no upsample in up_blocks.3
  _lowerCAmelCase      						=	F"""up_blocks.{i}.upsamplers.0."""
  _lowerCAmelCase      						=	F"""output_blocks.{3*i + 2}.{1 if i == 0 else 2}."""
  unet_conversion_map_layer.append((sd_upsample_prefix, hf_upsample_prefix))

_lowerCAmelCase      						=	'''mid_block.attentions.0.'''
_lowerCAmelCase      						=	'''middle_block.1.'''
unet_conversion_map_layer.append((sd_mid_atn_prefix, hf_mid_atn_prefix))

for j in range(2):
 _lowerCAmelCase      						=	F"""mid_block.resnets.{j}."""
 _lowerCAmelCase      						=	F"""middle_block.{2*j}."""
 unet_conversion_map_layer.append((sd_mid_res_prefix, hf_mid_res_prefix))
def    _SCREAMING_SNAKE_CASE  (       UpperCamelCase			):



    """simple docstring"""





    lowerCAmelCase__ :  Any							      =					{k: k for k in unet_state_dict.keys()}
    for sd_name, hf_name in unet_conversion_map:
        lowerCAmelCase__ :  Optional[int]							      =					sd_name
    for k, v in mapping.items():
        if "resnets" in k:
            for sd_part, hf_part in unet_conversion_map_resnet:
                lowerCAmelCase__ :  Any							      =					v.replace(UpperCamelCase							,    UpperCamelCase			)
            lowerCAmelCase__ :  List[Any]							      =					v
    for k, v in mapping.items():
        for sd_part, hf_part in unet_conversion_map_layer:
            lowerCAmelCase__ :  List[Any]							      =					v.replace(UpperCamelCase							,    UpperCamelCase			)
        lowerCAmelCase__ :  Optional[Any]							      =					v
    lowerCAmelCase__ :  Tuple							      =					{v: unet_state_dict[k] for k, v in mapping.items()}
    return new_state_dict


# ================#
# VAE Conversion #
# ================#

_lowerCAmelCase      						=	[
    # (stable-diffusion, HF Diffusers)
    ('''nin_shortcut''', '''conv_shortcut'''),
    ('''norm_out''', '''conv_norm_out'''),
    ('''mid.attn_1.''', '''mid_block.attentions.0.'''),
]

for i in range(4):
 # down_blocks have two resnets
 for j in range(2):
  _lowerCAmelCase      						=	F"""encoder.down_blocks.{i}.resnets.{j}."""
  _lowerCAmelCase      						=	F"""encoder.down.{i}.block.{j}."""
  vae_conversion_map.append((sd_down_prefix, hf_down_prefix))

 if i < 3:
  _lowerCAmelCase      						=	F"""down_blocks.{i}.downsamplers.0."""
  _lowerCAmelCase      						=	F"""down.{i}.downsample."""
  vae_conversion_map.append((sd_downsample_prefix, hf_downsample_prefix))

  _lowerCAmelCase      						=	F"""up_blocks.{i}.upsamplers.0."""
  _lowerCAmelCase      						=	F"""up.{3-i}.upsample."""
  vae_conversion_map.append((sd_upsample_prefix, hf_upsample_prefix))

 # up_blocks have three resnets
 # also, up blocks in hf are numbered in reverse from sd
 for j in range(3):
  _lowerCAmelCase      						=	F"""decoder.up_blocks.{i}.resnets.{j}."""
  _lowerCAmelCase      						=	F"""decoder.up.{3-i}.block.{j}."""
  vae_conversion_map.append((sd_up_prefix, hf_up_prefix))

# this part accounts for mid blocks in both the encoder and the decoder
for i in range(2):
 _lowerCAmelCase      						=	F"""mid_block.resnets.{i}."""
 _lowerCAmelCase      						=	F"""mid.block_{i+1}."""
 vae_conversion_map.append((sd_mid_res_prefix, hf_mid_res_prefix))


_lowerCAmelCase      						=	[
    # (stable-diffusion, HF Diffusers)
    ('''norm.''', '''group_norm.'''),
    ('''q.''', '''query.'''),
    ('''k.''', '''key.'''),
    ('''v.''', '''value.'''),
    ('''proj_out.''', '''proj_attn.'''),
]
def    _SCREAMING_SNAKE_CASE  (       UpperCamelCase			):



    """simple docstring"""





    return w.reshape(*w.shape							,    1							,    1			)
def    _SCREAMING_SNAKE_CASE  (       UpperCamelCase			):



    """simple docstring"""





    lowerCAmelCase__ :  Optional[int]							      =					{k: k for k in vae_state_dict.keys()}
    for k, v in mapping.items():
        for sd_part, hf_part in vae_conversion_map:
            lowerCAmelCase__ :  str							      =					v.replace(UpperCamelCase							,    UpperCamelCase			)
        lowerCAmelCase__ :  Optional[Any]							      =					v
    for k, v in mapping.items():
        if "attentions" in k:
            for sd_part, hf_part in vae_conversion_map_attn:
                lowerCAmelCase__ :  Dict							      =					v.replace(UpperCamelCase							,    UpperCamelCase			)
            lowerCAmelCase__ :  List[Any]							      =					v
    lowerCAmelCase__ :  Union[str, Any]							      =					{v: vae_state_dict[k] for k, v in mapping.items()}
    lowerCAmelCase__ :  Tuple							      =					["""q""", """k""", """v""", """proj_out"""]
    for k, v in new_state_dict.items():
        for weight_name in weights_to_convert:
            if f"""mid.attn_1.{weight_name}.weight""" in k:
                print(f"""Reshaping {k} for SD format"""			)
                lowerCAmelCase__ :  Optional[int]							      =					reshape_weight_for_sd(UpperCamelCase			)
    return new_state_dict


# =========================#
# Text Encoder Conversion #
# =========================#


_lowerCAmelCase      						=	[
    # (stable-diffusion, HF Diffusers)
    ('''resblocks.''', '''text_model.encoder.layers.'''),
    ('''ln_1''', '''layer_norm1'''),
    ('''ln_2''', '''layer_norm2'''),
    ('''.c_fc.''', '''.fc1.'''),
    ('''.c_proj.''', '''.fc2.'''),
    ('''.attn''', '''.self_attn'''),
    ('''ln_final.''', '''transformer.text_model.final_layer_norm.'''),
    ('''token_embedding.weight''', '''transformer.text_model.embeddings.token_embedding.weight'''),
    ('''positional_embedding''', '''transformer.text_model.embeddings.position_embedding.weight'''),
]
_lowerCAmelCase      						=	{re.escape(x[1]): x[0] for x in textenc_conversion_lst}
_lowerCAmelCase      						=	re.compile('''|'''.join(protected.keys()))

# Ordering is from https://github.com/pytorch/pytorch/blob/master/test/cpp/api/modules.cpp
_lowerCAmelCase      						=	{'''q''': 0, '''k''': 1, '''v''': 2}
def    _SCREAMING_SNAKE_CASE  (       UpperCamelCase			):



    """simple docstring"""





    lowerCAmelCase__ :  Optional[Any]							      =					{}
    lowerCAmelCase__ :  int							      =					{}
    lowerCAmelCase__ :  List[Any]							      =					{}
    for k, v in text_enc_dict.items():
        if (
            k.endswith(""".self_attn.q_proj.weight"""			)
            or k.endswith(""".self_attn.k_proj.weight"""			)
            or k.endswith(""".self_attn.v_proj.weight"""			)
        ):
            lowerCAmelCase__ :  Optional[int]							      =					k[: -len(""".q_proj.weight"""			)]
            lowerCAmelCase__ :  Tuple							      =					k[-len("""q_proj.weight"""			)]
            if k_pre not in capture_qkv_weight:
                lowerCAmelCase__ :  List[Any]							      =					[None, None, None]
            lowerCAmelCase__ :  Dict							      =					v
            continue

        if (
            k.endswith(""".self_attn.q_proj.bias"""			)
            or k.endswith(""".self_attn.k_proj.bias"""			)
            or k.endswith(""".self_attn.v_proj.bias"""			)
        ):
            lowerCAmelCase__ :  str							      =					k[: -len(""".q_proj.bias"""			)]
            lowerCAmelCase__ :  List[str]							      =					k[-len("""q_proj.bias"""			)]
            if k_pre not in capture_qkv_bias:
                lowerCAmelCase__ :  Union[str, Any]							      =					[None, None, None]
            lowerCAmelCase__ :  Any							      =					v
            continue

        lowerCAmelCase__ :  Dict							      =					textenc_pattern.sub(lambda UpperCamelCase			: protected[re.escape(m.group(0			)			)]							,    UpperCamelCase			)
        lowerCAmelCase__ :  Union[str, Any]							      =					v

    for k_pre, tensors in capture_qkv_weight.items():
        if None in tensors:
            raise Exception("""CORRUPTED MODEL: one of the q-k-v values for the text encoder was missing"""			)
        lowerCAmelCase__ :  Any							      =					textenc_pattern.sub(lambda UpperCamelCase			: protected[re.escape(m.group(0			)			)]							,    UpperCamelCase			)
        lowerCAmelCase__ :  Tuple							      =					torch.cat(UpperCamelCase			)

    for k_pre, tensors in capture_qkv_bias.items():
        if None in tensors:
            raise Exception("""CORRUPTED MODEL: one of the q-k-v values for the text encoder was missing"""			)
        lowerCAmelCase__ :  str							      =					textenc_pattern.sub(lambda UpperCamelCase			: protected[re.escape(m.group(0			)			)]							,    UpperCamelCase			)
        lowerCAmelCase__ :  List[Any]							      =					torch.cat(UpperCamelCase			)

    return new_state_dict
def    _SCREAMING_SNAKE_CASE  (       UpperCamelCase			):



    """simple docstring"""





    return text_enc_dict


if __name__ == "__main__":
 _lowerCAmelCase      						=	argparse.ArgumentParser()

 parser.add_argument('''--model_path''', default=None, type=str, required=True, help='''Path to the model to convert.''')
 parser.add_argument('''--checkpoint_path''', default=None, type=str, required=True, help='''Path to the output model.''')
 parser.add_argument('''--half''', action='''store_true''', help='''Save weights in half precision.''')
 parser.add_argument(
     '''--use_safetensors''', action='''store_true''', help='''Save weights use safetensors, default is ckpt.'''
 )

 _lowerCAmelCase      						=	parser.parse_args()

 assert args.model_path is not None, "Must provide a model path!"

 assert args.checkpoint_path is not None, "Must provide a checkpoint path!"

 # Path for safetensors
 _lowerCAmelCase      						=	osp.join(args.model_path, '''unet''', '''diffusion_pytorch_model.safetensors''')
 _lowerCAmelCase      						=	osp.join(args.model_path, '''vae''', '''diffusion_pytorch_model.safetensors''')
 _lowerCAmelCase      						=	osp.join(args.model_path, '''text_encoder''', '''model.safetensors''')

 # Load models from safetensors if it exists, if it doesn't pytorch
 if osp.exists(unet_path):
  _lowerCAmelCase      						=	load_file(unet_path, device='''cpu''')
 else:
  _lowerCAmelCase      						=	osp.join(args.model_path, '''unet''', '''diffusion_pytorch_model.bin''')
  _lowerCAmelCase      						=	torch.load(unet_path, map_location='''cpu''')

 if osp.exists(vae_path):
  _lowerCAmelCase      						=	load_file(vae_path, device='''cpu''')
 else:
  _lowerCAmelCase      						=	osp.join(args.model_path, '''vae''', '''diffusion_pytorch_model.bin''')
  _lowerCAmelCase      						=	torch.load(vae_path, map_location='''cpu''')

 if osp.exists(text_enc_path):
  _lowerCAmelCase      						=	load_file(text_enc_path, device='''cpu''')
 else:
  _lowerCAmelCase      						=	osp.join(args.model_path, '''text_encoder''', '''pytorch_model.bin''')
  _lowerCAmelCase      						=	torch.load(text_enc_path, map_location='''cpu''')

 # Convert the UNet model
 _lowerCAmelCase      						=	convert_unet_state_dict(unet_state_dict)
 _lowerCAmelCase      						=	{'''model.diffusion_model.''' + k: v for k, v in unet_state_dict.items()}

 # Convert the VAE model
 _lowerCAmelCase      						=	convert_vae_state_dict(vae_state_dict)
 _lowerCAmelCase      						=	{'''first_stage_model.''' + k: v for k, v in vae_state_dict.items()}

 # Easiest way to identify v2.0 model seems to be that the text encoder (OpenCLIP) is deeper
 _lowerCAmelCase      						=	'''text_model.encoder.layers.22.layer_norm2.bias''' in text_enc_dict

 if is_vaa_model:
  # Need to add the tag 'transformer' in advance so we can knock it out from the final layer-norm
  _lowerCAmelCase      						=	{'''transformer.''' + k: v for k, v in text_enc_dict.items()}
  _lowerCAmelCase      						=	convert_text_enc_state_dict_vaa(text_enc_dict)
  _lowerCAmelCase      						=	{'''cond_stage_model.model.''' + k: v for k, v in text_enc_dict.items()}
 else:
  _lowerCAmelCase      						=	convert_text_enc_state_dict(text_enc_dict)
  _lowerCAmelCase      						=	{'''cond_stage_model.transformer.''' + k: v for k, v in text_enc_dict.items()}

 # Put together new checkpoint
 _lowerCAmelCase      						=	{**unet_state_dict, **vae_state_dict, **text_enc_dict}
 if args.half:
  _lowerCAmelCase      						=	{k: v.half() for k, v in state_dict.items()}

 if args.use_safetensors:
  save_file(state_dict, args.checkpoint_path)
 else:
  _lowerCAmelCase      						=	{'''state_dict''': state_dict}
  torch.save(state_dict, args.checkpoint_path)


