



"""simple docstring"""






import gc
import unittest

from diffusers import FlaxDPMSolverMultistepScheduler, FlaxStableDiffusionPipeline
from diffusers.utils import is_flax_available, slow
from diffusers.utils.testing_utils import require_flax


if is_flax_available():
   import jax
   import jax.numpy as jnp
   from flax.jax_utils import replicate
   from flax.training.common_utils import shard






@slow
@require_flax
class 	__snake_case	(     unittest.TestCase):




  def 						SCREAMING_SNAKE_CASE    (			self					:						Any		):






         """simple docstring"""






         super().tearDown()
         gc.collect()




  def 						SCREAMING_SNAKE_CASE    (			self					:						Dict		):






         """simple docstring"""






         _lowerCamelCase      ,   _lowerCamelCase						:    Any         =			FlaxStableDiffusionPipeline.from_pretrained(
             '''stabilityai/stable-diffusion-2'''						, revision='''bf16'''						, dtype=jnp.bfloataa						, )

         _lowerCamelCase						:    Tuple         =			'''A painting of a squirrel eating a burger'''
         _lowerCamelCase						:    Dict         =			jax.device_count()
         _lowerCamelCase						:    Optional[Any]         =			num_samples * [prompt]
         _lowerCamelCase						:    Any         =			sd_pipe.prepare_inputs(__lowerCAmelCase		)

         _lowerCamelCase						:    Any         =			replicate(__lowerCAmelCase		)
         _lowerCamelCase						:    Tuple         =			shard(__lowerCAmelCase		)

         _lowerCamelCase						:    List[str]         =			jax.random.PRNGKey(0		)
         _lowerCamelCase						:    Tuple         =			jax.random.split(__lowerCAmelCase						, jax.device_count()		)

         _lowerCamelCase						:    str         =			sd_pipe(__lowerCAmelCase						, __lowerCAmelCase						, __lowerCAmelCase						, num_inference_steps=2_5						, jit=__lowerCAmelCase		)[0]
         assert images.shape == (jax.device_count(), 1, 7_6_8, 7_6_8, 3)

         _lowerCamelCase						:    Tuple         =			images.reshape((images.shape[0] * images.shape[1],) + images.shape[-3:]		)
         _lowerCamelCase						:    Dict         =			images[0, 2_5_3:2_5_6, 2_5_3:2_5_6, -1]

         _lowerCamelCase						:    str         =			jnp.asarray(jax.device_get(image_slice.flatten()		)		)
         _lowerCamelCase						:    Optional[int]         =			jnp.array([0.42_38, 0.44_14, 0.43_95, 0.44_53, 0.46_29, 0.45_90, 0.45_31, 0.4_55_08, 0.45_12]		)
         print(f'''output_slice: {output_slice}'''		)
         assert jnp.abs(output_slice - expected_slice		).max() < 1E-2




  def 						SCREAMING_SNAKE_CASE    (			self					:						Optional[Any]		):






         """simple docstring"""






         _lowerCamelCase						:    List[Any]         =			'''stabilityai/stable-diffusion-2'''
         _lowerCamelCase      ,   _lowerCamelCase						:    Optional[int]         =			FlaxDPMSolverMultistepScheduler.from_pretrained(__lowerCAmelCase						, subfolder='''scheduler'''		)
         _lowerCamelCase      ,   _lowerCamelCase						:    Dict         =			FlaxStableDiffusionPipeline.from_pretrained(
             __lowerCAmelCase						, scheduler=__lowerCAmelCase						, revision='''bf16'''						, dtype=jnp.bfloataa						, )
         _lowerCamelCase						:    Optional[int]         =			scheduler_params

         _lowerCamelCase						:    Union[str, Any]         =			'''A painting of a squirrel eating a burger'''
         _lowerCamelCase						:    Any         =			jax.device_count()
         _lowerCamelCase						:    Tuple         =			num_samples * [prompt]
         _lowerCamelCase						:    str         =			sd_pipe.prepare_inputs(__lowerCAmelCase		)

         _lowerCamelCase						:    str         =			replicate(__lowerCAmelCase		)
         _lowerCamelCase						:    Any         =			shard(__lowerCAmelCase		)

         _lowerCamelCase						:    Union[str, Any]         =			jax.random.PRNGKey(0		)
         _lowerCamelCase						:    int         =			jax.random.split(__lowerCAmelCase						, jax.device_count()		)

         _lowerCamelCase						:    Tuple         =			sd_pipe(__lowerCAmelCase						, __lowerCAmelCase						, __lowerCAmelCase						, num_inference_steps=2_5						, jit=__lowerCAmelCase		)[0]
         assert images.shape == (jax.device_count(), 1, 7_6_8, 7_6_8, 3)

         _lowerCamelCase						:    Any         =			images.reshape((images.shape[0] * images.shape[1],) + images.shape[-3:]		)
         _lowerCamelCase						:    int         =			images[0, 2_5_3:2_5_6, 2_5_3:2_5_6, -1]

         _lowerCamelCase						:    Dict         =			jnp.asarray(jax.device_get(image_slice.flatten()		)		)
         _lowerCamelCase						:    List[str]         =			jnp.array([0.43_36, 0.4_29_69, 0.44_53, 0.41_99, 0.42_97, 0.45_31, 0.44_34, 0.44_34, 0.42_97]		)
         print(f'''output_slice: {output_slice}'''		)
         assert jnp.abs(output_slice - expected_slice		).max() < 1E-2







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

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

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






def 	snake_case_  (				A_ :		int,    A_ :		Dict						):


       '''simple docstring'''





       _lowerCamelCase						:    List[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
       _lowerCamelCase						:    Union[str, Any]         =			int(re.match(R'''.*layer_(\d*).*''',    A_						)[1]						)
       layer_number -= 3
       return F'''h.{layer_number}.''' + key






def 	snake_case_  (				A_ :		List[Any]						):


       '''simple docstring'''





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






def 	snake_case_  (				A_ :		str,    A_ :		Any,    A_ :		int,    A_ :		List[str],    A_ :		Any						):


       '''simple docstring'''





       if bloom_config_file == "":
              _lowerCamelCase						:    Dict         =			BloomConfig()
       else:
              _lowerCamelCase						:    Any         =			BloomConfig.from_json_file(A_						)

       if shard_model:
              _lowerCamelCase						:    Optional[int]         =			os.listdir(A_						)
              _lowerCamelCase						:    List[str]         =			sorted(filter(lambda A_						: s.startswith('''layer'''						) and "model_00" in s,    A_						)						)

              _lowerCamelCase						:    str         =			{'''weight_map''': {}, '''metadata''': {}}
              _lowerCamelCase						:    List[str]         =			0

              _lowerCamelCase						:    str         =			None

              _lowerCamelCase						:    str         =			BloomConfig()

              for j, file in enumerate(A_						):
                     print('''Processing file: {}'''.format(A_						)						)
                     _lowerCamelCase						:    List[Any]         =			None

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

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

                            if tensors is None:
                                   _lowerCamelCase						:    Any         =			temp
                            else:
                                   for key in tensors.keys():
                                          if any(key.endswith(A_						) 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						:    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
                                                 _lowerCamelCase						:    Optional[Any]         =			torch.cat([tensors[key], temp[key]],    dim=A_						)

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

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

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

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

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

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

                            if tensors is None:
                                   _lowerCamelCase						:    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(A_						) 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						:    Optional[Any]         =			torch.cat([tensors[key], temp[key]],    dim=A_						)

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

                     _lowerCamelCase						:    List[str]         =			model.load_state_dict(A_,    strict=A_						)
                     assert not other_keys.unexpected_keys, F'''The keys {other_keys.unexpected_keys} are unexpected'''
                     if missing_keys is None:
                            _lowerCamelCase						:    Optional[Any]         =			set(other_keys.missing_keys						)
                     else:
                            _lowerCamelCase						:    Union[str, 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(A_,    exist_ok=A_						)
              _lowerCamelCase						:    Union[str, Any]         =			pytorch_dump_folder_path + '''/''' + WEIGHTS_NAME
              _lowerCamelCase						:    Tuple         =			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						:    Dict         =			model.to(config.torch_dtype						)
              torch.save(model.state_dict(),    A_						)
              print(F'''Save configuration file to {pytorch_config_dump_path}'''						)
              with open(A_,    '''w''',    encoding='''utf-8'''						) as f:
                     f.write(config.to_json_string()						)


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


