






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      :      Dict             =       [
    """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      :      str             =       [
    """mlp.dense_4h_to_h.weight""",
    """self_attention.dense.weight""",
]







def 		_UpperCamelCase				(      UpperCamelCase_    :   str      ,    UpperCamelCase_    :   Dict						)					->				Any:

						"""simple docstring"""



						lowerCAmelCase__        =  {
						    '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__        =  int(re.match(r'.*layer_(\d*).*'      ,    UpperCamelCase_						)[1]						)
						layer_number -= 3
						return F"h.{layer_number}." + key







def 		_UpperCamelCase				(      UpperCamelCase_    :   Union[str, Any]						)					->				List[str]:

						"""simple docstring"""



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







def 		_UpperCamelCase				(      UpperCamelCase_    :   Union[str, Any]      ,    UpperCamelCase_    :   Tuple      ,    UpperCamelCase_    :   Tuple      ,    UpperCamelCase_    :   Tuple      ,    UpperCamelCase_    :   List[Any]						)					->				List[Any]:

						"""simple docstring"""



						if bloom_config_file == "":
												lowerCAmelCase__        =  BloomConfig()
						else:
												lowerCAmelCase__        =  BloomConfig.from_json_file(UpperCamelCase_						)

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

												lowerCAmelCase__        =  {'weight_map': {}, 'metadata': {}}
												lowerCAmelCase__        =  0

												lowerCAmelCase__        =  None

												lowerCAmelCase__        =  BloomConfig()

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

																		for i in range(UpperCamelCase_						):
																								# load all TP files
																								lowerCAmelCase__        =  file.replace('model_00'      ,    F"model_0{i}"						)
																								lowerCAmelCase__        =  torch.load(os.path.join(UpperCamelCase_      ,    UpperCamelCase_						)      ,    map_location='cpu'						)

																								# Rename keys in the transformers names
																								lowerCAmelCase__        =  list(temp.keys()						)
																								for key in keys:
																														lowerCAmelCase__        =  temp.pop(UpperCamelCase_						)

																								if tensors is None:
																														lowerCAmelCase__        =  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
																																										lowerCAmelCase__        =  1 if any(text in key for text in WEIGHTS_WITH_ROW_PARALLELISM_CONTAIN						) else 0
																																										# We concatenate these weights accross TP ranks
																																										lowerCAmelCase__        =  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						):
																														lowerCAmelCase__        =  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():
																								lowerCAmelCase__        =  tensors[key]
																								total_size += value.numel() * get_dtype_size(value.dtype						)
																								if key not in index_dict["weight_map"]:
																														lowerCAmelCase__        =  'pytorch_model_{}-of-{}.bin'.format(
																														    str(j + 1						).zfill(5						)      ,    str(len(UpperCamelCase_						)						).zfill(5						)						)

												lowerCAmelCase__        =  BloomConfig()
												lowerCAmelCase__        =  pytorch_dump_folder_path + '/' + CONFIG_NAME
												lowerCAmelCase__        =  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:
																		lowerCAmelCase__        =  json.dumps(UpperCamelCase_      ,    indent=2      ,    sort_keys=UpperCamelCase_						) + '\n'
																		f.write(UpperCamelCase_						)
						else:
												lowerCAmelCase__        =  BloomModel(UpperCamelCase_						)

												lowerCAmelCase__        =  os.listdir(UpperCamelCase_						)
												lowerCAmelCase__        =  sorted(filter(lambda UpperCamelCase_						: s.startswith('layer'						) and "model_00" in s      ,    UpperCamelCase_						)						)

												lowerCAmelCase__        =  None
												for i, file in enumerate(UpperCamelCase_						):
																		lowerCAmelCase__        =  None
																		for i in range(UpperCamelCase_						):
																								# load all TP files
																								lowerCAmelCase__        =  file.replace('model_00'      ,    F"model_0{i}"						)
																								lowerCAmelCase__        =  torch.load(os.path.join(UpperCamelCase_      ,    UpperCamelCase_						)      ,    map_location='cpu'						)

																								# Rename keys in the transformers names
																								lowerCAmelCase__        =  list(temp.keys()						)
																								for key in keys:
																														lowerCAmelCase__        =  temp.pop(UpperCamelCase_						)

																								if tensors is None:
																														lowerCAmelCase__        =  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
																																										lowerCAmelCase__        =  1 if any(text in key for text in WEIGHTS_WITH_ROW_PARALLELISM_CONTAIN						) else 0
																																										# We concatenate these weights accross TP ranks
																																										lowerCAmelCase__        =  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						):
																														lowerCAmelCase__        =  tensors[key] / pretraining_tp

																		lowerCAmelCase__        =  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:
																								lowerCAmelCase__        =  set(other_keys.missing_keys						)
																		else:
																								lowerCAmelCase__        =  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_						)
												lowerCAmelCase__        =  pytorch_dump_folder_path + '/' + WEIGHTS_NAME
												lowerCAmelCase__        =  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__        =  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__":
					__snake_case      :      Optional[Any]             =       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      :      str             =       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 unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    DPMSolverMultistepScheduler,
    TextToVideoSDPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import is_xformers_available, load_numpy, skip_mps, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineTesterMixin


enable_full_determinism()


@skip_mps
class       __SCREAMING_SNAKE_CASE   (				__lowercase						,   unittest.TestCase):
						_SCREAMING_SNAKE_CASE		:    List[str]						=   TextToVideoSDPipeline
						_SCREAMING_SNAKE_CASE		:    str						=   TEXT_TO_IMAGE_PARAMS
						_SCREAMING_SNAKE_CASE		:    List[str]						=   TEXT_TO_IMAGE_BATCH_PARAMS
						# No `output_type`.
						_SCREAMING_SNAKE_CASE		:    Dict						=   frozenset(
						    [
						        '''num_inference_steps''',
						        '''generator''',
						        '''latents''',
						        '''return_dict''',
						        '''callback''',
						        '''callback_steps''',
						    ])






						def  UpperCamelCase__ (		self				):






												"""simple docstring"""
												torch.manual_seed(0				)
												lowerCAmelCase__        =  UNetaDConditionModel(
												    block_out_channels=(32, 64, 64, 64)       ,  layers_per_block=2       ,  sample_size=32       ,  in_channels=4       ,  out_channels=4       ,  down_block_types=('CrossAttnDownBlock3D', 'CrossAttnDownBlock3D', 'CrossAttnDownBlock3D', 'DownBlock3D')       ,  up_block_types=('UpBlock3D', 'CrossAttnUpBlock3D', 'CrossAttnUpBlock3D', 'CrossAttnUpBlock3D')       ,  cross_attention_dim=32       ,  attention_head_dim=4       ,  )
												lowerCAmelCase__        =  DDIMScheduler(
												    beta_start=0.0_00_85       ,  beta_end=0.0_12       ,  beta_schedule='scaled_linear'       ,  clip_sample=_UpperCamelCase       ,  set_alpha_to_one=_UpperCamelCase       ,  )
												torch.manual_seed(0				)
												lowerCAmelCase__        =  AutoencoderKL(
												    block_out_channels=[32, 64]       ,  in_channels=3       ,  out_channels=3       ,  down_block_types=['DownEncoderBlock2D', 'DownEncoderBlock2D']       ,  up_block_types=['UpDecoderBlock2D', 'UpDecoderBlock2D']       ,  latent_channels=4       ,  sample_size=1_28       ,  )
												torch.manual_seed(0				)
												lowerCAmelCase__        =  CLIPTextConfig(
												    bos_token_id=0       ,  eos_token_id=2       ,  hidden_size=32       ,  intermediate_size=37       ,  layer_norm_eps=1E-05       ,  num_attention_heads=4       ,  num_hidden_layers=5       ,  pad_token_id=1       ,  vocab_size=10_00       ,  hidden_act='gelu'       ,  projection_dim=5_12       ,  )
												lowerCAmelCase__        =  CLIPTextModel(_UpperCamelCase				)
												lowerCAmelCase__        =  CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'				)

												lowerCAmelCase__        =  {
												    'unet': unet,
												    'scheduler': scheduler,
												    'vae': vae,
												    'text_encoder': text_encoder,
												    'tokenizer': tokenizer,
												}
												return components






						def  UpperCamelCase__ (		self       ,  _UpperCamelCase       ,  _UpperCamelCase=0				):






												"""simple docstring"""
												if str(_UpperCamelCase				).startswith('mps'				):
																		lowerCAmelCase__        =  torch.manual_seed(_UpperCamelCase				)
												else:
																		lowerCAmelCase__        =  torch.Generator(device=_UpperCamelCase				).manual_seed(_UpperCamelCase				)
												lowerCAmelCase__        =  {
												    'prompt': 'A painting of a squirrel eating a burger',
												    'generator': generator,
												    'num_inference_steps': 2,
												    'guidance_scale': 6.0,
												    'output_type': 'pt',
												}
												return inputs






						def  UpperCamelCase__ (		self				):






												"""simple docstring"""
												lowerCAmelCase__        =  'cpu'  # ensure determinism for the device-dependent torch.Generator
												lowerCAmelCase__        =  self.get_dummy_components()
												lowerCAmelCase__        =  TextToVideoSDPipeline(**_UpperCamelCase				)
												lowerCAmelCase__        =  sd_pipe.to(_UpperCamelCase				)
												sd_pipe.set_progress_bar_config(disable=_UpperCamelCase				)

												lowerCAmelCase__        =  self.get_dummy_inputs(_UpperCamelCase				)
												lowerCAmelCase__        =  'np'
												lowerCAmelCase__        =  sd_pipe(**_UpperCamelCase				).frames
												lowerCAmelCase__        =  frames[0][-3:, -3:, -1]

												assert frames[0].shape == (64, 64, 3)
												lowerCAmelCase__        =  np.array([1_58.0, 1_60.0, 1_53.0, 1_25.0, 1_00.0, 1_21.0, 1_11.0, 93.0, 1_13.0]				)

												assert np.abs(image_slice.flatten() - expected_slice				).max() < 1E-2






						def  UpperCamelCase__ (		self				):






												"""simple docstring"""
												self._test_attention_slicing_forward_pass(test_mean_pixel_difference=_UpperCamelCase       ,  expected_max_diff=3E-3				)






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






												"""simple docstring"""
												self._test_xformers_attention_forwardGenerator_pass(test_mean_pixel_difference=_UpperCamelCase       ,  expected_max_diff=1E-2				)






						@unittest.skip(reason='Batching needs to be properly figured out first for this pipeline.'				)
						def  UpperCamelCase__ (		self				):






												"""simple docstring"""
												pass






						@unittest.skip(reason='Batching needs to be properly figured out first for this pipeline.'				)
						def  UpperCamelCase__ (		self				):






												"""simple docstring"""
												pass






						@unittest.skip(reason='`num_images_per_prompt` argument is not supported for this pipeline.'				)
						def  UpperCamelCase__ (		self				):






												"""simple docstring"""
												pass







						def  UpperCamelCase__ (		self				):






												"""simple docstring"""
												return super().test_progress_bar()




@slow
@skip_mps
class       __SCREAMING_SNAKE_CASE   (				unittest.TestCase):






						def  UpperCamelCase__ (		self				):






												"""simple docstring"""
												lowerCAmelCase__        =  load_numpy(
												    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/text_to_video/video.npy'				)

												lowerCAmelCase__        =  TextToVideoSDPipeline.from_pretrained('damo-vilab/text-to-video-ms-1.7b'				)
												lowerCAmelCase__        =  DPMSolverMultistepScheduler.from_config(pipe.scheduler.config				)
												lowerCAmelCase__        =  pipe.to('cuda'				)

												lowerCAmelCase__        =  'Spiderman is surfing'
												lowerCAmelCase__        =  torch.Generator(device='cpu'				).manual_seed(0				)

												lowerCAmelCase__        =  pipe(_UpperCamelCase       ,  generator=_UpperCamelCase       ,  num_inference_steps=25       ,  output_type='pt'				).frames
												lowerCAmelCase__        =  video_frames.cpu().numpy()

												assert np.abs(expected_video - video				).mean() < 5E-2







						def  UpperCamelCase__ (		self				):






												"""simple docstring"""
												lowerCAmelCase__        =  load_numpy(
												    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/text_to_video/video_2step.npy'				)

												lowerCAmelCase__        =  TextToVideoSDPipeline.from_pretrained('damo-vilab/text-to-video-ms-1.7b'				)
												lowerCAmelCase__        =  pipe.to('cuda'				)

												lowerCAmelCase__        =  'Spiderman is surfing'
												lowerCAmelCase__        =  torch.Generator(device='cpu'				).manual_seed(0				)

												lowerCAmelCase__        =  pipe(_UpperCamelCase       ,  generator=_UpperCamelCase       ,  num_inference_steps=2       ,  output_type='pt'				).frames
												lowerCAmelCase__        =  video_frames.cpu().numpy()

												assert np.abs(expected_video - video				).mean() < 5E-2





