





'''simple docstring'''

import gc
import random
import unittest

import numpy as np
import torch

from diffusers import DDIMScheduler, KandinskyVaaPipeline, KandinskyVaaPriorPipeline, UNetaDConditionModel, VQModel
from diffusers.utils import floats_tensor, load_numpy, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..test_pipelines_common import PipelineTesterMixin, assert_mean_pixel_difference


enable_full_determinism()
class 							__lowerCAmelCase	(       __magic_name__		,		unittest.TestCase					):






       """simple docstring"""
       _snake_case      :	str									=       KandinskyVaaPipeline
       _snake_case      :	Union[str, Any]									=       [
           'image_embeds',
           'negative_image_embeds',
       ]
       _snake_case      :	Optional[Any]									=       ['image_embeds', 'negative_image_embeds']
       _snake_case      :	Union[str, Any]									=       [
           'generator',
           'height',
           'width',
           'latents',
           'guidance_scale',
           'num_inference_steps',
           'return_dict',
           'guidance_scale',
           'num_images_per_prompt',
           'output_type',
           'return_dict',
       ]
       _snake_case      :	List[Any]									=       False





       @property
       def snake_case__		(				self  :						Tuple	)       ->	List[Any]:

            '''simple docstring'''




            return 32





       @property
       def snake_case__		(				self  :						int	)       ->	List[str]:

            '''simple docstring'''




            return 32





       @property
       def snake_case__		(				self  :						List[str]	)       ->	List[Any]:

            '''simple docstring'''




            return self.time_input_dim





       @property
       def snake_case__		(				self  :						List[str]	)       ->	Optional[Any]:

            '''simple docstring'''




            return self.time_input_dim * 4





       @property
       def snake_case__		(				self  :						Tuple	)       ->	List[str]:

            '''simple docstring'''




            return 100





       @property
       def snake_case__		(				self  :						Optional[Any]	)       ->	Optional[Any]:

            '''simple docstring'''




            torch.manual_seed(0	)

            _UpperCamelCase							       =			{
                '''in_channels''': 4,
                # Out channels is double in channels because predicts mean and variance
                '''out_channels''': 8,
                '''addition_embed_type''': '''image''',
                '''down_block_types''': ('''ResnetDownsampleBlock2D''', '''SimpleCrossAttnDownBlock2D'''),
                '''up_block_types''': ('''SimpleCrossAttnUpBlock2D''', '''ResnetUpsampleBlock2D'''),
                '''mid_block_type''': '''UNetMidBlock2DSimpleCrossAttn''',
                '''block_out_channels''': (self.block_out_channels_a, self.block_out_channels_a * 2),
                '''layers_per_block''': 1,
                '''encoder_hid_dim''': self.text_embedder_hidden_size,
                '''encoder_hid_dim_type''': '''image_proj''',
                '''cross_attention_dim''': self.cross_attention_dim,
                '''attention_head_dim''': 4,
                '''resnet_time_scale_shift''': '''scale_shift''',
                '''class_embed_type''': None,
            }

            _UpperCamelCase							       =			UNetaDConditionModel(**lowerCAmelCase__	)
            return model





       @property
       def snake_case__		(				self  :						Optional[int]	)       ->	Tuple:

            '''simple docstring'''




            return {
                "block_out_channels": [32, 64],
                "down_block_types": ["DownEncoderBlock2D", "AttnDownEncoderBlock2D"],
                "in_channels": 3,
                "latent_channels": 4,
                "layers_per_block": 1,
                "norm_num_groups": 8,
                "norm_type": "spatial",
                "num_vq_embeddings": 12,
                "out_channels": 3,
                "up_block_types": [
                    "AttnUpDecoderBlock2D",
                    "UpDecoderBlock2D",
                ],
                "vq_embed_dim": 4,
            }





       @property
       def snake_case__		(				self  :						Dict	)       ->	List[str]:

            '''simple docstring'''




            torch.manual_seed(0	)
            _UpperCamelCase							       =			VQModel(**self.dummy_movq_kwargs	)
            return model





       def snake_case__		(				self  :						Dict	)       ->	str:

            '''simple docstring'''




            _UpperCamelCase							       =			self.dummy_unet
            _UpperCamelCase							       =			self.dummy_movq

            _UpperCamelCase							       =			DDIMScheduler(
                num_train_timesteps=1000			,					beta_schedule='''linear'''			,					beta_start=0.00085			,					beta_end=0.012			,					clip_sample=lowerCAmelCase__			,					set_alpha_to_one=lowerCAmelCase__			,					steps_offset=1			,					prediction_type='''epsilon'''			,					thresholding=lowerCAmelCase__			,					)

            _UpperCamelCase							       =			{
                '''unet''': unet,
                '''scheduler''': scheduler,
                '''movq''': movq,
            }
            return components





       def snake_case__		(				self  :						List[str]			,					lowerCAmelCase__  :						List[str]			,					lowerCAmelCase__  :						Tuple=0	)       ->	str:

            '''simple docstring'''




            _UpperCamelCase							       =			floats_tensor((1, self.text_embedder_hidden_size)			,					rng=random.Random(lowerCAmelCase__	)	).to(lowerCAmelCase__	)
            _UpperCamelCase							       =			floats_tensor((1, self.text_embedder_hidden_size)			,					rng=random.Random(seed + 1	)	).to(
                lowerCAmelCase__	)
            if str(lowerCAmelCase__	).startswith('''mps'''	):
                 _UpperCamelCase							       =			torch.manual_seed(lowerCAmelCase__	)
            else:
                 _UpperCamelCase							       =			torch.Generator(device=lowerCAmelCase__	).manual_seed(lowerCAmelCase__	)
            _UpperCamelCase							       =			{
                '''image_embeds''': image_embeds,
                '''negative_image_embeds''': negative_image_embeds,
                '''generator''': generator,
                '''height''': 64,
                '''width''': 64,
                '''guidance_scale''': 4.0,
                '''num_inference_steps''': 2,
                '''output_type''': '''np''',
            }
            return inputs







       def snake_case__		(				self  :						Any	)       ->	Dict:

            '''simple docstring'''




            _UpperCamelCase							       =			'''cpu'''

            _UpperCamelCase							       =			self.get_dummy_components()

            _UpperCamelCase							       =			self.pipeline_class(**lowerCAmelCase__	)
            _UpperCamelCase							       =			pipe.to(lowerCAmelCase__	)

            pipe.set_progress_bar_config(disable=lowerCAmelCase__	)

            _UpperCamelCase							       =			pipe(**self.get_dummy_inputs(lowerCAmelCase__	)	)
            _UpperCamelCase							       =			output.images

            _UpperCamelCase							       =			pipe(
                **self.get_dummy_inputs(lowerCAmelCase__	)			,					return_dict=lowerCAmelCase__			,					)[0]

            _UpperCamelCase							       =			image[0, -3:, -3:, -1]
            _UpperCamelCase							       =			image_from_tuple[0, -3:, -3:, -1]

            assert image.shape == (1, 64, 64, 3)

            _UpperCamelCase							       =			np.array(
                [0.6237976, 1.0, 0.36441332, 1.0, 0.70639634, 0.29877186, 0.85652125, 0.5216843, 0.54454046]	)

            assert (
                np.abs(image_slice.flatten() - expected_slice	).max() < 1e-2
            ), f""" expected_slice {expected_slice}, but got {image_slice.flatten()}"""

            assert (
                np.abs(image_from_tuple_slice.flatten() - expected_slice	).max() < 1e-2
            ), f""" expected_slice {expected_slice}, but got {image_from_tuple_slice.flatten()}"""

@slow
@require_torch_gpu
class 							__lowerCAmelCase	(       unittest.TestCase					):






       """simple docstring"""





       def snake_case__		(				self  :						Any	)       ->	Any:

            '''simple docstring'''




            super().tearDown()
            gc.collect()
            torch.cuda.empty_cache()







       def snake_case__		(				self  :						List[Any]	)       ->	List[Any]:

            '''simple docstring'''




            _UpperCamelCase							       =			load_numpy(
                '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
                '''/kandinskyv22/kandinskyv22_text2img_cat_fp16.npy'''	)

            _UpperCamelCase							       =			KandinskyVaaPriorPipeline.from_pretrained(
                '''kandinsky-community/kandinsky-2-2-prior'''			,					torch_dtype=torch.floataa	)
            pipe_prior.to(lowerCAmelCase__	)

            _UpperCamelCase							       =			KandinskyVaaPipeline.from_pretrained(
                '''kandinsky-community/kandinsky-2-2-decoder'''			,					torch_dtype=torch.floataa	)
            _UpperCamelCase							       =			pipeline.to(lowerCAmelCase__	)
            pipeline.set_progress_bar_config(disable=lowerCAmelCase__	)

            _UpperCamelCase							       =			'''red cat, 4k photo'''

            _UpperCamelCase							       =			torch.Generator(device='''cuda'''	).manual_seed(0	)
            _UpperCamelCase       ,							_UpperCamelCase							       =			pipe_prior(
                lowerCAmelCase__			,					generator=lowerCAmelCase__			,					num_inference_steps=5			,					negative_prompt=''''''			,					).to_tuple()

            _UpperCamelCase							       =			torch.Generator(device='''cuda'''	).manual_seed(0	)
            _UpperCamelCase							       =			pipeline(
                image_embeds=lowerCAmelCase__			,					negative_image_embeds=lowerCAmelCase__			,					generator=lowerCAmelCase__			,					num_inference_steps=100			,					output_type='''np'''			,					)

            _UpperCamelCase							       =			output.images[0]

            assert image.shape == (512, 512, 3)

            assert_mean_pixel_difference(lowerCAmelCase__			,					lowerCAmelCase__	)








'''simple docstring'''

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()

a							   =				[
    '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',
]

a							   =				[
    'mlp.dense_4h_to_h.weight',
    'self_attention.dense.weight',
]







def a_     (    __UpperCAmelCase				,							__UpperCAmelCase					)				->					Any:





			"""simple docstring"""







			snake_case:    Any      							={
			    '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
			snake_case:    Dict      							=int(re.match(R'.*layer_(\d*).*'				,							__UpperCAmelCase					)[1]					)
			layer_number -= 3
			return f'''h.{layer_number}.''' + key







def a_     (    __UpperCAmelCase					)				->					Union[str, Any]:





			"""simple docstring"""







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







def a_     (    __UpperCAmelCase				,							__UpperCAmelCase				,							__UpperCAmelCase				,							__UpperCAmelCase				,							__UpperCAmelCase					)				->					Tuple:





			"""simple docstring"""







			if bloom_config_file == "":
						snake_case:    Tuple      							=BloomConfig()
			else:
						snake_case:    List[Any]      							=BloomConfig.from_json_file(__UpperCAmelCase					)

			if shard_model:
						snake_case:    Any      							=os.listdir(__UpperCAmelCase					)
						snake_case:    Any      							=sorted(filter(lambda __UpperCAmelCase					: s.startswith('layer'					) and "model_00" in s				,							__UpperCAmelCase					)					)

						snake_case:    int      							={'weight_map': {}, 'metadata': {}}
						snake_case:    Dict      							=0

						snake_case:    List[str]      							=None

						snake_case:    Any      							=BloomConfig()

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

									for i in range(__UpperCAmelCase					):
												# load all TP files
												snake_case:    Optional[int]      							=file.replace('model_00'				,							f'''model_0{i}'''					)
												snake_case:    List[str]      							=torch.load(os.path.join(__UpperCAmelCase				,							__UpperCAmelCase					)				,							map_location='cpu'					)

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

												if tensors is None:
															snake_case:    Tuple      							=temp
												else:
															for key in tensors.keys():
																		if any(key.endswith(__UpperCAmelCase					) 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
																					snake_case:    Optional[int]      							=1 if any(text in key for text in WEIGHTS_WITH_ROW_PARALLELISM_CONTAIN					) else 0
																					# We concatenate these weights accross TP ranks
																					snake_case:    str      							=torch.cat([tensors[key], temp[key]]				,							dim=__UpperCAmelCase					)

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

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

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

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

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

												# Rename keys in the transformers names
												snake_case:    List[str]      							=list(temp.keys()					)
												for key in keys:
															snake_case:    int      							=temp.pop(__UpperCAmelCase					)

												if tensors is None:
															snake_case:    Any      							=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(__UpperCAmelCase					) for end in WEIGHTS_TO_AVERAGE_ENDSWITH					):
																					tensors[key] += temp[key]
																		else:
																					# Some weights are RowParallelLinear in Megatron-Deepspeed, others are ColumnParallel
																					snake_case:    List[Any]      							=1 if any(text in key for text in WEIGHTS_WITH_ROW_PARALLELISM_CONTAIN					) else 0
																					# We concatenate these weights accross TP ranks
																					snake_case:    Tuple      							=torch.cat([tensors[key], temp[key]]				,							dim=__UpperCAmelCase					)

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

									snake_case:    Optional[Any]      							=model.load_state_dict(__UpperCAmelCase				,							strict=__UpperCAmelCase					)
									assert not other_keys.unexpected_keys, f'''The keys {other_keys.unexpected_keys} are unexpected'''
									if missing_keys is None:
												snake_case:    Optional[Any]      							=set(other_keys.missing_keys					)
									else:
												snake_case:    Any      							=missing_keys.intersection(set(other_keys.missing_keys					)					)

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

						# Save pytorch-model
						os.makedirs(__UpperCAmelCase				,							exist_ok=__UpperCAmelCase					)
						snake_case:    Optional[int]      							=pytorch_dump_folder_path + '/' + WEIGHTS_NAME
						snake_case:    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:
									snake_case:    Optional[Any]      							=model.to(config.torch_dtype					)
						torch.save(model.state_dict()				,							__UpperCAmelCase					)
						print(f'''Save configuration file to {pytorch_config_dump_path}'''					)
						with open(__UpperCAmelCase				,							'w'				,							encoding='utf-8'					) as f:
									f.write(config.to_json_string()					)


if __name__ == "__main__":
	a							   =				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',
	)
	a							   =				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,
	)




