



import argparse
import json
import os
import re

import torch

from transformers import BloomConfig, BloomModel
from transformers.file_utils import CONFIG_NAME, WEIGHTS_NAME
from transformers.utils import logging


logging.set_verbosity_info()

__snake_case    :		List[str]  =[
    'word_embeddings_layernorm.weight',
    'word_embeddings_layernorm.bias',
    'input_layernorm.weight',
    'input_layernorm.bias',
    'post_attention_layernorm.weight',
    'post_attention_layernorm.bias',
    'self_attention.dense.bias',
    'mlp.dense_4h_to_h.bias',
    'ln_f.weight',
    'ln_f.bias',
]

__snake_case    :		Optional[int]  =[
    'mlp.dense_4h_to_h.weight',
    'self_attention.dense.weight',
]






def      lowerCAmelCase__		(   lowerCamelCase_		:		int	,lowerCamelCase_		:		List[Any]):





				'''simple docstring'''

				lowerCAmelCase__					:    Tuple												=					{
				    '''word_embeddings.weight''': '''word_embeddings.weight''',
				    '''word_embeddings.norm.weight''': '''word_embeddings_layernorm.weight''',
				    '''word_embeddings.norm.bias''': '''word_embeddings_layernorm.bias''',
				    '''weight''': '''ln_f.weight''',
				    '''bias''': '''ln_f.bias''',
				}

				if key in layer_rename_map:
								return layer_rename_map[key]

				# Handle transformer blocks
				lowerCAmelCase__					:    Any												=					int(re.match(r'''.*layer_(\d*).*'''	,lowerCamelCase_)[1])
				layer_number -= 3
				return f"""h.{layer_number}.""" + key






def      lowerCAmelCase__		(   lowerCamelCase_		:		List[str]):





				'''simple docstring'''

				if dtype == torch.bool:
								return 1 / 8
				lowerCAmelCase__					:    Tuple												=					re.search(r'''[^\d](\d+)$'''	,str(lowerCamelCase_))
				if bit_search is None:
								raise ValueError(f"""`dtype` is not a valid dtype: {dtype}.""")
				lowerCAmelCase__					:    Tuple												=					int(bit_search.groups()[0])
				return bit_size // 8






def      lowerCAmelCase__		(   lowerCamelCase_		:		Optional[Any]	,lowerCamelCase_		:		Tuple	,lowerCamelCase_		:		int	,lowerCamelCase_		:		Dict	,lowerCamelCase_		:		int):





				'''simple docstring'''

				if bloom_config_file == "":
								lowerCAmelCase__					:    Optional[Any]												=					BloomConfig()
				else:
								lowerCAmelCase__					:    Tuple												=					BloomConfig.from_json_file(lowerCamelCase_)

				if shard_model:
								lowerCAmelCase__					:    List[Any]												=					os.listdir(lowerCamelCase_)
								lowerCAmelCase__					:    str												=					sorted(filter(lambda lowerCamelCase_: s.startswith('''layer''') and "model_00" in s	,lowerCamelCase_))

								lowerCAmelCase__					:    Optional[int]												=					{'''weight_map''': {}, '''metadata''': {}}
								lowerCAmelCase__					:    Optional[Any]												=					0

								lowerCAmelCase__					:    str												=					None

								lowerCAmelCase__					:    int												=					BloomConfig()

								for j, file in enumerate(lowerCamelCase_):
												print('''Processing file: {}'''.format(lowerCamelCase_))
												lowerCAmelCase__					:    Any												=					None

												for i in range(lowerCamelCase_):
																# load all TP files
																lowerCAmelCase__					:    List[str]												=					file.replace('''model_00'''	,f"""model_0{i}""")
																lowerCAmelCase__					:    Union[str, Any]												=					torch.load(os.path.join(lowerCamelCase_	,lowerCamelCase_)	,map_location='''cpu''')

																# Rename keys in the transformers names
																lowerCAmelCase__					:    List[Any]												=					list(temp.keys())
																for key in keys:
																				lowerCAmelCase__					:    List[Any]												=					temp.pop(lowerCamelCase_)

																if tensors is None:
																				lowerCAmelCase__					:    Optional[int]												=					temp
																else:
																				for key in tensors.keys():
																								if any(key.endswith(lowerCamelCase_) for end in WEIGHTS_TO_AVERAGE_ENDSWITH):
																												# We average (sum and then divide) some weights accross TP ranks (see https://github.com/bigscience-workshop/Megatron-DeepSpeed/blob/olruwase/sync_layer_norms/megatron/training.py#L425)
																												tensors[key] += temp[key]
																								else:
																												# Some weights are RowParallelLinear in Megatron-Deepspeed, others are ColumnParallel
																												lowerCAmelCase__					:    Union[str, Any]												=					1 if any(text in key for text in WEIGHTS_WITH_ROW_PARALLELISM_CONTAIN) else 0
																												# We concatenate these weights accross TP ranks
																												lowerCAmelCase__					:    Tuple												=					torch.cat([tensors[key], temp[key]]	,dim=lowerCamelCase_)

            # Divide by the number of TP the weights we want to average
												for key in tensors.keys():
																if any(key.endswith(lowerCamelCase_) for end in WEIGHTS_TO_AVERAGE_ENDSWITH):
																				lowerCAmelCase__					:    Optional[int]												=					tensors[key] / pretraining_tp
												torch.save(
												    lowerCamelCase_	,os.path.join(
												        lowerCamelCase_	,'''pytorch_model_{}-of-{}.bin'''.format(str(j + 1).zfill(5)	,str(len(lowerCamelCase_)).zfill(5))	,)	,)

												for key in tensors.keys():
																lowerCAmelCase__					:    Union[str, Any]												=					tensors[key]
																total_size += value.numel() * get_dtype_size(value.dtype)
																if key not in index_dict["weight_map"]:
																				lowerCAmelCase__					:    List[str]												=					'''pytorch_model_{}-of-{}.bin'''.format(
																				    str(j + 1).zfill(5)	,str(len(lowerCamelCase_)).zfill(5))

								lowerCAmelCase__					:    Optional[int]												=					BloomConfig()
								lowerCAmelCase__					:    Union[str, Any]												=					pytorch_dump_folder_path + '''/''' + CONFIG_NAME
								lowerCAmelCase__					:    List[Any]												=					total_size
								with open(lowerCamelCase_	,'''w'''	,encoding='''utf-8''') as f:
												f.write(config.to_json_string())
								with open(os.path.join(lowerCamelCase_	,WEIGHTS_NAME + '''.index.json''')	,'''w'''	,encoding='''utf-8''') as f:
												lowerCAmelCase__					:    Tuple												=					json.dumps(lowerCamelCase_	,indent=2	,sort_keys=lowerCamelCase_) + '''\n'''
												f.write(lowerCamelCase_)
				else:
								lowerCAmelCase__					:    str												=					BloomModel(lowerCamelCase_)

								lowerCAmelCase__					:    Union[str, Any]												=					os.listdir(lowerCamelCase_)
								lowerCAmelCase__					:    Optional[int]												=					sorted(filter(lambda lowerCamelCase_: s.startswith('''layer''') and "model_00" in s	,lowerCamelCase_))

								lowerCAmelCase__					:    List[str]												=					None
								for i, file in enumerate(lowerCamelCase_):
												lowerCAmelCase__					:    int												=					None
												for i in range(lowerCamelCase_):
																# load all TP files
																lowerCAmelCase__					:    List[str]												=					file.replace('''model_00'''	,f"""model_0{i}""")
																lowerCAmelCase__					:    List[Any]												=					torch.load(os.path.join(lowerCamelCase_	,lowerCamelCase_)	,map_location='''cpu''')

																# Rename keys in the transformers names
																lowerCAmelCase__					:    Union[str, Any]												=					list(temp.keys())
																for key in keys:
																				lowerCAmelCase__					:    int												=					temp.pop(lowerCamelCase_)

																if tensors is None:
																				lowerCAmelCase__					:    Optional[int]												=					temp
																else:
																				for key in tensors.keys():
																								# We average (sum and then divide) some weights accross TP ranks (see https://github.com/bigscience-workshop/Megatron-DeepSpeed/blob/olruwase/sync_layer_norms/megatron/training.py#L425)
																								if any(key.endswith(lowerCamelCase_) for end in WEIGHTS_TO_AVERAGE_ENDSWITH):
																												tensors[key] += temp[key]
																								else:
																												# Some weights are RowParallelLinear in Megatron-Deepspeed, others are ColumnParallel
																												lowerCAmelCase__					:    Dict												=					1 if any(text in key for text in WEIGHTS_WITH_ROW_PARALLELISM_CONTAIN) else 0
																												# We concatenate these weights accross TP ranks
																												lowerCAmelCase__					:    Union[str, Any]												=					torch.cat([tensors[key], temp[key]]	,dim=lowerCamelCase_)

            # Divide by the number of TP the weights we want to average
												for key in tensors.keys():
																if any(key.endswith(lowerCamelCase_) for end in WEIGHTS_TO_AVERAGE_ENDSWITH):
																				lowerCAmelCase__					:    Optional[int]												=					tensors[key] / pretraining_tp

												lowerCAmelCase__					:    Any												=					model.load_state_dict(lowerCamelCase_	,strict=lowerCamelCase_)
												assert not other_keys.unexpected_keys, f"""The keys {other_keys.unexpected_keys} are unexpected"""
												if missing_keys is None:
																lowerCAmelCase__					:    List[str]												=					set(other_keys.missing_keys)
												else:
																lowerCAmelCase__					:    Optional[int]												=					missing_keys.intersection(set(other_keys.missing_keys))

								assert not missing_keys, f"""The keys {missing_keys} are missing"""

								# Save pytorch-model
								os.makedirs(lowerCamelCase_	,exist_ok=lowerCamelCase_)
								lowerCAmelCase__					:    Optional[int]												=					pytorch_dump_folder_path + '''/''' + WEIGHTS_NAME
								lowerCAmelCase__					:    Any												=					pytorch_dump_folder_path + '''/''' + CONFIG_NAME
								print(f"""Save PyTorch model to {pytorch_weights_dump_path} with dtype {config.torch_dtype}""")
								if config.torch_dtype is not None:
												lowerCAmelCase__					:    Any												=					model.to(config.torch_dtype)
								torch.save(model.state_dict()	,lowerCamelCase_)
								print(f"""Save configuration file to {pytorch_config_dump_path}""")
								with open(lowerCamelCase_	,'''w'''	,encoding='''utf-8''') as f:
												f.write(config.to_json_string())


if __name__ == "__main__":
							__snake_case    :		Dict  =argparse.ArgumentParser()
							# Required parameters
							parser.add_argument(
							    '--bloom_checkpoint_path',
							    default=None,
							    type=str,
							    required=True,
							    help='Path to the Megatron-LM checkpoint path.',
							)
							parser.add_argument(
							    '--pytorch_dump_folder_path', default=None, type=str, required=True, help='Path to the output PyTorch model.'
							)
							parser.add_argument(
							    '--bloom_config_file',
							    default='',
							    type=str,
							    help=(
							        'An optional config json file corresponding to the pre-trained model. \n'
							        'This specifies the model architecture.'
							    ),
							)
							parser.add_argument(
							    '--shard_model',
							    action='store_true',
							    help='An optional setting to shard the output model \nThis enables sharding the converted checkpoint',
							)
							parser.add_argument(
							    '--pretraining_tp',
							    default=4,
							    type=int,
							    help='Pretraining TP rank that has been used when training the model in Megatron-LM \n',
							)
							__snake_case    :		Union[str, Any]  =parser.parse_args()
							convert_bloom_checkpoint_to_pytorch(
							    args.bloom_checkpoint_path,
							    args.bloom_config_file,
							    args.pytorch_dump_folder_path,
							    args.shard_model,
							    args.pretraining_tp,
							)






import random
import unittest

import torch

from diffusers import IFInpaintingSuperResolutionPipeline
from diffusers.utils import floats_tensor
from diffusers.utils.import_utils import is_xformers_available
from diffusers.utils.testing_utils import skip_mps, torch_device

from ..pipeline_params import (
    TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS,
    TEXT_GUIDED_IMAGE_INPAINTING_PARAMS,
)
from ..test_pipelines_common import PipelineTesterMixin
from . import IFPipelineTesterMixin



@skip_mps
class 	lowerCamelCase__      (							lowerCamelCase__				,    lowerCamelCase__				,    unittest.TestCase):




		'''simple docstring'''



		snake_case_     =IFInpaintingSuperResolutionPipeline
		snake_case_     =TEXT_GUIDED_IMAGE_INPAINTING_PARAMS - {"""width""", """height"""}
		snake_case_     =TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS.union({"""original_image"""})
		snake_case_     =PipelineTesterMixin.required_optional_params - {"""latents"""}



		def 		lowerCAmelCase__				(self			)					->      Union[str, Any]:

						"""simple docstring"""






						return self._get_superresolution_dummy_components()



		def 		lowerCAmelCase__				(self		,__lowerCamelCase		,__lowerCamelCase=0			)					->      Dict:

						"""simple docstring"""






						if str(__lowerCamelCase			).startswith('''mps'''			):
										lowerCAmelCase__					:    Union[str, Any]												=					torch.manual_seed(__lowerCamelCase			)
						else:
										lowerCAmelCase__					:    int												=					torch.Generator(device=__lowerCamelCase			).manual_seed(__lowerCamelCase			)

						lowerCAmelCase__					:    str												=					floats_tensor((1, 3, 16, 16)		,rng=random.Random(__lowerCamelCase			)			).to(__lowerCamelCase			)
						lowerCAmelCase__					:    Optional[Any]												=					floats_tensor((1, 3, 32, 32)		,rng=random.Random(__lowerCamelCase			)			).to(__lowerCamelCase			)
						lowerCAmelCase__					:    Union[str, Any]												=					floats_tensor((1, 3, 32, 32)		,rng=random.Random(__lowerCamelCase			)			).to(__lowerCamelCase			)

						lowerCAmelCase__					:    Dict												=					{
						    '''prompt''': '''A painting of a squirrel eating a burger''',
						    '''image''': image,
						    '''original_image''': original_image,
						    '''mask_image''': mask_image,
						    '''generator''': generator,
						    '''num_inference_steps''': 2,
						    '''output_type''': '''numpy''',
						}

						return inputs



		@unittest.skipIf(
		    torch_device != '''cuda''' or not is_xformers_available()		,reason='''XFormers attention is only available with CUDA and `xformers` installed'''		,)
		def 		lowerCAmelCase__				(self			)					->      int:

						"""simple docstring"""






						self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=1e-3			)



		def 		lowerCAmelCase__				(self			)					->      int:

						"""simple docstring"""






						self._test_save_load_optional_components()



		@unittest.skipIf(torch_device != '''cuda'''		,reason='''float16 requires CUDA'''			)
		def 		lowerCAmelCase__				(self			)					->      str:

						"""simple docstring"""






						super().test_save_load_floataa(expected_max_diff=1e-1			)



		def 		lowerCAmelCase__				(self			)					->      Any:

						"""simple docstring"""






						self._test_attention_slicing_forward_pass(expected_max_diff=1e-2			)



		def 		lowerCAmelCase__				(self			)					->      Optional[Any]:

						"""simple docstring"""






						self._test_save_load_local()



		def 		lowerCAmelCase__				(self			)					->      Any:

						"""simple docstring"""






						self._test_inference_batch_single_identical(
						    expected_max_diff=1e-2		,)

