



"""simple docstring"""






import gc
import random
import tempfile
import unittest

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

from diffusers import (
    AutoencoderKL,
    DDIMInverseScheduler,
    DDIMScheduler,
    DPMSolverMultistepInverseScheduler,
    DPMSolverMultistepScheduler,
    StableDiffusionDiffEditPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import load_image, slow
from diffusers.utils.testing_utils import enable_full_determinism, floats_tensor, require_torch_gpu, torch_device

from ..pipeline_params import TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS, TEXT_GUIDED_IMAGE_INPAINTING_PARAMS
from ..test_pipelines_common import PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()






class 	__snake_case	(     _lowercase  ,					_lowercase  ,					unittest.TestCase):
  snake_case__    :       Any        =						StableDiffusionDiffEditPipeline
  snake_case__    :       Optional[Any]        =						TEXT_GUIDED_IMAGE_INPAINTING_PARAMS - {"height", "width", "image"} | {"image_latents"}
  snake_case__    :       Tuple        =						TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS - {"image"} | {"image_latents"}
  snake_case__    :       Any        =						frozenset(
      [])  # TO-DO: update image_params once pipeline is refactored with VaeImageProcessor.preprocess
  snake_case__    :       int        =						frozenset([])




  def 						SCREAMING_SNAKE_CASE    (			self					:						str		):






         """simple docstring"""






         torch.manual_seed(0		)
         _lowerCamelCase						:    Optional[Any]         =			UNetaDConditionModel(
             block_out_channels=(3_2, 6_4)						, layers_per_block=2						, sample_size=3_2						, in_channels=4						, out_channels=4						, down_block_types=('''DownBlock2D''', '''CrossAttnDownBlock2D''')						, up_block_types=('''CrossAttnUpBlock2D''', '''UpBlock2D''')						, cross_attention_dim=3_2						, attention_head_dim=(2, 4)						, use_linear_projection=__lowerCAmelCase						, )
         _lowerCamelCase						:    List[str]         =			DDIMScheduler(
             beta_start=0.0_00_85						, beta_end=0.0_12						, beta_schedule='''scaled_linear'''						, clip_sample=__lowerCAmelCase						, set_alpha_to_one=__lowerCAmelCase						, )
         _lowerCamelCase						:    int         =			DDIMInverseScheduler(
             beta_start=0.0_00_85						, beta_end=0.0_12						, beta_schedule='''scaled_linear'''						, clip_sample=__lowerCAmelCase						, set_alpha_to_zero=__lowerCAmelCase						, )
         torch.manual_seed(0		)
         _lowerCamelCase						:    List[str]         =			AutoencoderKL(
             block_out_channels=[3_2, 6_4]						, in_channels=3						, out_channels=3						, down_block_types=['''DownEncoderBlock2D''', '''DownEncoderBlock2D''']						, up_block_types=['''UpDecoderBlock2D''', '''UpDecoderBlock2D''']						, latent_channels=4						, sample_size=1_2_8						, )
         torch.manual_seed(0		)
         _lowerCamelCase						:    int         =			CLIPTextConfig(
             bos_token_id=0						, eos_token_id=2						, hidden_size=3_2						, intermediate_size=3_7						, layer_norm_eps=1E-05						, num_attention_heads=4						, num_hidden_layers=5						, pad_token_id=1						, vocab_size=1_0_0_0						, hidden_act='''gelu'''						, projection_dim=5_1_2						, )
         _lowerCamelCase						:    Tuple         =			CLIPTextModel(__lowerCAmelCase		)
         _lowerCamelCase						:    Dict         =			CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''		)

         _lowerCamelCase						:    List[Any]         =			{
             '''unet''': unet,
             '''scheduler''': scheduler,
             '''inverse_scheduler''': inverse_scheduler,
             '''vae''': vae,
             '''text_encoder''': text_encoder,
             '''tokenizer''': tokenizer,
             '''safety_checker''': None,
             '''feature_extractor''': None,
         }

         return components




  def 						SCREAMING_SNAKE_CASE    (			self					:						str						, __lowerCAmelCase					:						str						, __lowerCAmelCase					:						List[Any]=0		):






         """simple docstring"""






         _lowerCamelCase						:    Optional[int]         =			floats_tensor((1, 1_6, 1_6)						, rng=random.Random(__lowerCAmelCase		)		).to(__lowerCAmelCase		)
         _lowerCamelCase						:    Union[str, Any]         =			floats_tensor((1, 2, 4, 1_6, 1_6)						, rng=random.Random(__lowerCAmelCase		)		).to(__lowerCAmelCase		)
         if str(__lowerCAmelCase		).startswith('''mps'''		):
                _lowerCamelCase						:    str         =			torch.manual_seed(__lowerCAmelCase		)
         else:
                _lowerCamelCase						:    Dict         =			torch.Generator(device=__lowerCAmelCase		).manual_seed(__lowerCAmelCase		)
         _lowerCamelCase						:    Optional[Any]         =			{
             '''prompt''': '''a dog and a newt''',
             '''mask_image''': mask,
             '''image_latents''': latents,
             '''generator''': generator,
             '''num_inference_steps''': 2,
             '''inpaint_strength''': 1.0,
             '''guidance_scale''': 6.0,
             '''output_type''': '''numpy''',
         }

         return inputs




  def 						SCREAMING_SNAKE_CASE    (			self					:						List[Any]						, __lowerCAmelCase					:						Tuple						, __lowerCAmelCase					:						List[str]=0		):






         """simple docstring"""






         _lowerCamelCase						:    Union[str, Any]         =			floats_tensor((1, 3, 3_2, 3_2)						, rng=random.Random(__lowerCAmelCase		)		).to(__lowerCAmelCase		)
         _lowerCamelCase						:    Optional[int]         =			image.cpu().permute(0						, 2						, 3						, 1		)[0]
         _lowerCamelCase						:    Union[str, Any]         =			Image.fromarray(np.uinta(__lowerCAmelCase		)		).convert('''RGB'''		)
         if str(__lowerCAmelCase		).startswith('''mps'''		):
                _lowerCamelCase						:    List[str]         =			torch.manual_seed(__lowerCAmelCase		)
         else:
                _lowerCamelCase						:    str         =			torch.Generator(device=__lowerCAmelCase		).manual_seed(__lowerCAmelCase		)
         _lowerCamelCase						:    Optional[Any]         =			{
             '''image''': image,
             '''source_prompt''': '''a cat and a frog''',
             '''target_prompt''': '''a dog and a newt''',
             '''generator''': generator,
             '''num_inference_steps''': 2,
             '''num_maps_per_mask''': 2,
             '''mask_encode_strength''': 1.0,
             '''guidance_scale''': 6.0,
             '''output_type''': '''numpy''',
         }

         return inputs




  def 						SCREAMING_SNAKE_CASE    (			self					:						List[str]						, __lowerCAmelCase					:						Optional[int]						, __lowerCAmelCase					:						List[str]=0		):






         """simple docstring"""






         _lowerCamelCase						:    Union[str, Any]         =			floats_tensor((1, 3, 3_2, 3_2)						, rng=random.Random(__lowerCAmelCase		)		).to(__lowerCAmelCase		)
         _lowerCamelCase						:    List[Any]         =			image.cpu().permute(0						, 2						, 3						, 1		)[0]
         _lowerCamelCase						:    Optional[int]         =			Image.fromarray(np.uinta(__lowerCAmelCase		)		).convert('''RGB'''		)
         if str(__lowerCAmelCase		).startswith('''mps'''		):
                _lowerCamelCase						:    Union[str, Any]         =			torch.manual_seed(__lowerCAmelCase		)
         else:
                _lowerCamelCase						:    Any         =			torch.Generator(device=__lowerCAmelCase		).manual_seed(__lowerCAmelCase		)
         _lowerCamelCase						:    Tuple         =			{
             '''image''': image,
             '''prompt''': '''a cat and a frog''',
             '''generator''': generator,
             '''num_inference_steps''': 2,
             '''inpaint_strength''': 1.0,
             '''guidance_scale''': 6.0,
             '''decode_latents''': True,
             '''output_type''': '''numpy''',
         }
         return inputs




  def 						SCREAMING_SNAKE_CASE    (			self					:						str		):






         """simple docstring"""






         if not hasattr(self.pipeline_class						, '''_optional_components'''		):
                return

         _lowerCamelCase						:    Dict         =			self.get_dummy_components()
         _lowerCamelCase						:    str         =			self.pipeline_class(**__lowerCAmelCase		)
         pipe.to(__lowerCAmelCase		)
         pipe.set_progress_bar_config(disable=__lowerCAmelCase		)

         # set all optional components to None and update pipeline config accordingly
         for optional_component in pipe._optional_components:
                setattr(__lowerCAmelCase						, __lowerCAmelCase						, __lowerCAmelCase		)
         pipe.register_modules(**{optional_component: None for optional_component in pipe._optional_components}		)

         _lowerCamelCase						:    int         =			self.get_dummy_inputs(__lowerCAmelCase		)
         _lowerCamelCase						:    Optional[int]         =			pipe(**__lowerCAmelCase		)[0]

         with tempfile.TemporaryDirectory() as tmpdir:
                pipe.save_pretrained(__lowerCAmelCase		)
                _lowerCamelCase						:    List[Any]         =			self.pipeline_class.from_pretrained(__lowerCAmelCase		)
                pipe_loaded.to(__lowerCAmelCase		)
                pipe_loaded.set_progress_bar_config(disable=__lowerCAmelCase		)

         for optional_component in pipe._optional_components:
                self.assertTrue(
                    getattr(__lowerCAmelCase						, __lowerCAmelCase		) is None						, f'''`{optional_component}` did not stay set to None after loading.'''						, )

         _lowerCamelCase						:    List[Any]         =			self.get_dummy_inputs(__lowerCAmelCase		)
         _lowerCamelCase						:    Any         =			pipe_loaded(**__lowerCAmelCase		)[0]

         _lowerCamelCase						:    int         =			np.abs(output - output_loaded		).max()
         self.assertLess(__lowerCAmelCase						, 1E-4		)




  def 						SCREAMING_SNAKE_CASE    (			self					:						Any		):






         """simple docstring"""






         _lowerCamelCase						:    List[Any]         =			'''cpu'''

         _lowerCamelCase						:    List[str]         =			self.get_dummy_components()
         _lowerCamelCase						:    Tuple         =			self.pipeline_class(**__lowerCAmelCase		)
         pipe.to(__lowerCAmelCase		)
         pipe.set_progress_bar_config(disable=__lowerCAmelCase		)

         _lowerCamelCase						:    Dict         =			self.get_dummy_mask_inputs(__lowerCAmelCase		)
         _lowerCamelCase						:    Dict         =			pipe.generate_mask(**__lowerCAmelCase		)
         _lowerCamelCase						:    Dict         =			mask[0, -3:, -3:]

         self.assertEqual(mask.shape						, (1, 1_6, 1_6)		)
         _lowerCamelCase						:    int         =			np.array([0] * 9		)
         _lowerCamelCase						:    Union[str, Any]         =			np.abs(mask_slice.flatten() - expected_slice		).max()
         self.assertLessEqual(__lowerCAmelCase						, 1E-3		)
         self.assertEqual(mask[0, -3, -4]						, 0		)




  def 						SCREAMING_SNAKE_CASE    (			self					:						Optional[Any]		):






         """simple docstring"""






         _lowerCamelCase						:    int         =			'''cpu'''

         _lowerCamelCase						:    Any         =			self.get_dummy_components()
         _lowerCamelCase						:    Dict         =			self.pipeline_class(**__lowerCAmelCase		)
         pipe.to(__lowerCAmelCase		)
         pipe.set_progress_bar_config(disable=__lowerCAmelCase		)

         _lowerCamelCase						:    int         =			self.get_dummy_inversion_inputs(__lowerCAmelCase		)
         _lowerCamelCase						:    Tuple         =			pipe.invert(**__lowerCAmelCase		).images
         _lowerCamelCase						:    Optional[int]         =			image[0, -1, -3:, -3:]

         self.assertEqual(image.shape						, (2, 3_2, 3_2, 3)		)
         _lowerCamelCase						:    str         =			np.array(
             [0.51_50, 0.51_34, 0.50_43, 0.53_76, 0.46_94, 0.5_10_50, 0.50_15, 0.44_07, 0.47_99]						, )
         _lowerCamelCase						:    Any         =			np.abs(image_slice.flatten() - expected_slice		).max()
         self.assertLessEqual(__lowerCAmelCase						, 1E-3		)




  def 						SCREAMING_SNAKE_CASE    (			self					:						Any		):






         """simple docstring"""






         super().test_inference_batch_single_identical(expected_max_diff=5E-3		)




  def 						SCREAMING_SNAKE_CASE    (			self					:						List[str]		):






         """simple docstring"""






         _lowerCamelCase						:    Optional[Any]         =			'''cpu'''

         _lowerCamelCase						:    Optional[int]         =			self.get_dummy_components()

         _lowerCamelCase						:    List[Any]         =			{'''beta_start''': 0.0_00_85, '''beta_end''': 0.0_12, '''beta_schedule''': '''scaled_linear'''}
         _lowerCamelCase						:    List[Any]         =			DPMSolverMultistepScheduler(**__lowerCAmelCase		)
         _lowerCamelCase						:    List[Any]         =			DPMSolverMultistepInverseScheduler(**__lowerCAmelCase		)

         _lowerCamelCase						:    Dict         =			self.pipeline_class(**__lowerCAmelCase		)
         pipe.to(__lowerCAmelCase		)
         pipe.set_progress_bar_config(disable=__lowerCAmelCase		)

         _lowerCamelCase						:    List[Any]         =			self.get_dummy_inversion_inputs(__lowerCAmelCase		)
         _lowerCamelCase						:    Optional[Any]         =			pipe.invert(**__lowerCAmelCase		).images
         _lowerCamelCase						:    Any         =			image[0, -1, -3:, -3:]

         self.assertEqual(image.shape						, (2, 3_2, 3_2, 3)		)
         _lowerCamelCase						:    Dict         =			np.array(
             [0.51_50, 0.51_34, 0.50_43, 0.53_76, 0.46_94, 0.5_10_50, 0.50_15, 0.44_07, 0.47_99]						, )
         _lowerCamelCase						:    str         =			np.abs(image_slice.flatten() - expected_slice		).max()
         self.assertLessEqual(__lowerCAmelCase						, 1E-3		)






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




  def 						SCREAMING_SNAKE_CASE    (			self					:						List[Any]		):






         """simple docstring"""






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




  @classmethod
  def 						SCREAMING_SNAKE_CASE    (			cls					:						str		):






         """simple docstring"""






         _lowerCamelCase						:    List[Any]         =			load_image(
             '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/diffedit/fruit.png'''		)

         _lowerCamelCase						:    Dict         =			raw_image.convert('''RGB'''		).resize((7_6_8, 7_6_8)		)

         _lowerCamelCase						:    Dict         =			raw_image




  def 						SCREAMING_SNAKE_CASE    (			self					:						Optional[Any]		):






         """simple docstring"""






         _lowerCamelCase						:    int         =			torch.manual_seed(0		)

         _lowerCamelCase						:    Optional[int]         =			StableDiffusionDiffEditPipeline.from_pretrained(
             '''stabilityai/stable-diffusion-2-1'''						, safety_checker=__lowerCAmelCase						, torch_dtype=torch.floataa		)
         _lowerCamelCase						:    Union[str, Any]         =			DDIMScheduler.from_config(pipe.scheduler.config		)
         _lowerCamelCase						:    Any         =			DDIMInverseScheduler.from_config(pipe.scheduler.config		)
         pipe.enable_model_cpu_offload()
         pipe.set_progress_bar_config(disable=__lowerCAmelCase		)

         _lowerCamelCase						:    Any         =			'''a bowl of fruit'''
         _lowerCamelCase						:    int         =			'''a bowl of pears'''

         _lowerCamelCase						:    List[Any]         =			pipe.generate_mask(
             image=self.raw_image						, source_prompt=__lowerCAmelCase						, target_prompt=__lowerCAmelCase						, generator=__lowerCAmelCase						, )

         _lowerCamelCase						:    Optional[int]         =			pipe.invert(
             prompt=__lowerCAmelCase						, image=self.raw_image						, inpaint_strength=0.7						, generator=__lowerCAmelCase		).latents

         _lowerCamelCase						:    str         =			pipe(
             prompt=__lowerCAmelCase						, mask_image=__lowerCAmelCase						, image_latents=__lowerCAmelCase						, generator=__lowerCAmelCase						, negative_prompt=__lowerCAmelCase						, inpaint_strength=0.7						, output_type='''numpy'''						, ).images[0]

         _lowerCamelCase						:    str         =			(
             np.array(
                 load_image(
                     '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
                     '''/diffedit/pears.png'''		).resize((7_6_8, 7_6_8)		)		)
             / 2_5_5
         )
         assert np.abs((expected_image - image).max()		) < 5E-1




  def 						SCREAMING_SNAKE_CASE    (			self					:						Tuple		):






         """simple docstring"""






         _lowerCamelCase						:    Optional[int]         =			torch.manual_seed(0		)

         _lowerCamelCase						:    str         =			StableDiffusionDiffEditPipeline.from_pretrained(
             '''stabilityai/stable-diffusion-2-1'''						, safety_checker=__lowerCAmelCase						, torch_dtype=torch.floataa		)
         _lowerCamelCase						:    Any         =			DPMSolverMultistepScheduler.from_config(pipe.scheduler.config		)
         _lowerCamelCase						:    Any         =			DPMSolverMultistepInverseScheduler.from_config(pipe.scheduler.config		)
         pipe.enable_model_cpu_offload()
         pipe.set_progress_bar_config(disable=__lowerCAmelCase		)

         _lowerCamelCase						:    Optional[Any]         =			'''a bowl of fruit'''
         _lowerCamelCase						:    Union[str, Any]         =			'''a bowl of pears'''

         _lowerCamelCase						:    int         =			pipe.generate_mask(
             image=self.raw_image						, source_prompt=__lowerCAmelCase						, target_prompt=__lowerCAmelCase						, generator=__lowerCAmelCase						, )

         _lowerCamelCase						:    int         =			pipe.invert(
             prompt=__lowerCAmelCase						, image=self.raw_image						, inpaint_strength=0.7						, generator=__lowerCAmelCase						, num_inference_steps=2_5						, ).latents

         _lowerCamelCase						:    Dict         =			pipe(
             prompt=__lowerCAmelCase						, mask_image=__lowerCAmelCase						, image_latents=__lowerCAmelCase						, generator=__lowerCAmelCase						, negative_prompt=__lowerCAmelCase						, inpaint_strength=0.7						, num_inference_steps=2_5						, output_type='''numpy'''						, ).images[0]

         _lowerCamelCase						:    List[str]         =			(
             np.array(
                 load_image(
                     '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
                     '''/diffedit/pears.png'''		).resize((7_6_8, 7_6_8)		)		)
             / 2_5_5
         )
         assert np.abs((expected_image - image).max()		) < 5E-1







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


