
'''simple docstring'''


import inspect
from typing import Callable, List, Optional, Union

import torch
from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer

from diffusers import DiffusionPipeline
from diffusers.models import AutoencoderKL, UNetaDConditionModel
from diffusers.pipelines.stable_diffusion import StableDiffusionPipelineOutput
from diffusers.pipelines.stable_diffusion.safety_checker import StableDiffusionSafetyChecker
from diffusers.schedulers import DDIMScheduler, LMSDiscreteScheduler, PNDMScheduler
from diffusers.utils import logging


snake_case_ :  Dict     =					logging.get_logger(__name__)  # pylint: disable=invalid-name


class 						lowercase__	(  snake_case_   ):






   '''simple docstring'''







   def __init__(			self ,   lowerCamelCase__ ,   lowerCamelCase__ ,   lowerCamelCase__ ,   lowerCamelCase__ ,   lowerCamelCase__ ,   lowerCamelCase__ ,   lowerCamelCase__ ,   ):

      '''simple docstring'''







      super().__init__()
      self.register_modules(
          vae=lowerCamelCase__ ,   text_encoder=lowerCamelCase__ ,   tokenizer=lowerCamelCase__ ,   unet=lowerCamelCase__ ,   scheduler=lowerCamelCase__ ,   safety_checker=lowerCamelCase__ ,   feature_extractor=lowerCamelCase__ ,   )







   def       UpperCAmelCase				(			self ,   lowerCamelCase__ = "auto"   ):

      '''simple docstring'''







      if slice_size == "auto":
         # half the attention head size is usually a good trade-off between
         # speed and memory
         UpperCamelCase       =  self.unet.config.attention_head_dim // 2
      self.unet.set_attention_slice(lowerCamelCase__   )







   def       UpperCAmelCase				(			self   ):

      '''simple docstring'''







      self.enable_attention_slicing(lowerCamelCase__   )







   @torch.no_grad()
   def __call__(			self ,   lowerCamelCase__ ,   lowerCamelCase__ = 5_1_2 ,   lowerCamelCase__ = 5_1_2 ,   lowerCamelCase__ = 5_0 ,   lowerCamelCase__ = 7.5 ,   lowerCamelCase__ = None ,   lowerCamelCase__ = 1 ,   lowerCamelCase__ = 0.0 ,   lowerCamelCase__ = None ,   lowerCamelCase__ = None ,   lowerCamelCase__ = "pil" ,   lowerCamelCase__ = True ,   lowerCamelCase__ = None ,   lowerCamelCase__ = 1 ,   lowerCamelCase__ = None ,   **lowerCamelCase__ ,   ):

      '''simple docstring'''







      if isinstance(lowerCamelCase__ ,   lowerCamelCase__   ):
         UpperCamelCase       =  1
      elif isinstance(lowerCamelCase__ ,   lowerCamelCase__   ):
         UpperCamelCase       =  len(lowerCamelCase__   )
      else:
         raise ValueError(f'`prompt` has to be of type `str` or `list` but is {type(lowerCamelCase__   )}'   )

      if height % 8 != 0 or width % 8 != 0:
         raise ValueError(f'`height` and `width` have to be divisible by 8 but are {height} and {width}.'   )

      if (callback_steps is None) or (
          callback_steps is not None and (not isinstance(lowerCamelCase__ ,   lowerCamelCase__   ) or callback_steps <= 0)
      ):
         raise ValueError(
             f'`callback_steps` has to be a positive integer but is {callback_steps} of type'
             f' {type(lowerCamelCase__   )}.'   )

      # get prompt text embeddings
      UpperCamelCase       =  self.tokenizer(
          lowerCamelCase__ ,   padding='''max_length''' ,   max_length=self.tokenizer.model_max_length ,   return_tensors='''pt''' ,   )
      UpperCamelCase       =  text_inputs.input_ids

      if text_input_ids.shape[-1] > self.tokenizer.model_max_length:
         UpperCamelCase       =  self.tokenizer.batch_decode(text_input_ids[:, self.tokenizer.model_max_length :]   )
         logger.warning(
             '''The following part of your input was truncated because CLIP can only handle sequences up to'''
             f' {self.tokenizer.model_max_length} tokens: {removed_text}'   )
         UpperCamelCase       =  text_input_ids[:, : self.tokenizer.model_max_length]

      if text_embeddings is None:
         UpperCamelCase       =  self.text_encoder(text_input_ids.to(self.device   )   )[0]

      # duplicate text embeddings for each generation per prompt, using mps friendly method
      UpperCamelCase     , UpperCamelCase     , UpperCamelCase       =  text_embeddings.shape
      UpperCamelCase       =  text_embeddings.repeat(1 ,   lowerCamelCase__ ,   1   )
      UpperCamelCase       =  text_embeddings.view(bs_embed * num_images_per_prompt ,   lowerCamelCase__ ,   -1   )

      # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
      # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
      # corresponds to doing no classifier free guidance.
      UpperCamelCase       =  guidance_scale > 1.0
      # get unconditional embeddings for classifier free guidance
      if do_classifier_free_guidance:
         UpperCamelCase       =  42
         if negative_prompt is None:
            UpperCamelCase       =  ['''''']
         elif type(lowerCamelCase__   ) is not type(lowerCamelCase__   ):
            raise TypeError(
                f'`negative_prompt` should be the same type to `prompt`, but got {type(lowerCamelCase__   )} !='
                f' {type(lowerCamelCase__   )}.'   )
         elif isinstance(lowerCamelCase__ ,   lowerCamelCase__   ):
            UpperCamelCase       =  [negative_prompt]
         elif batch_size != len(lowerCamelCase__   ):
            raise ValueError(
                f'`negative_prompt`: {negative_prompt} has batch size {len(lowerCamelCase__   )}, but `prompt`:'
                f' {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches'
                ''' the batch size of `prompt`.'''   )
         else:
            UpperCamelCase       =  negative_prompt

         UpperCamelCase       =  text_input_ids.shape[-1]
         UpperCamelCase       =  self.tokenizer(
             lowerCamelCase__ ,   padding='''max_length''' ,   max_length=lowerCamelCase__ ,   truncation=lowerCamelCase__ ,   return_tensors='''pt''' ,   )
         UpperCamelCase       =  self.text_encoder(uncond_input.input_ids.to(self.device   )   )[0]

         # duplicate unconditional embeddings for each generation per prompt, using mps friendly method
         UpperCamelCase       =  uncond_embeddings.shape[1]
         UpperCamelCase       =  uncond_embeddings.repeat(lowerCamelCase__ ,   lowerCamelCase__ ,   1   )
         UpperCamelCase       =  uncond_embeddings.view(batch_size * num_images_per_prompt ,   lowerCamelCase__ ,   -1   )

         # For classifier free guidance, we need to do two forward passes.
         # Here we concatenate the unconditional and text embeddings into a single batch
         # to avoid doing two forward passes
         UpperCamelCase       =  torch.cat([uncond_embeddings, text_embeddings]   )

      # get the initial random noise unless the user supplied it

      # Unlike in other pipelines, latents need to be generated in the target device
      # for 1-to-1 results reproducibility with the CompVis implementation.
      # However this currently doesn't work in `mps`.
      UpperCamelCase       =  (batch_size * num_images_per_prompt, self.unet.config.in_channels, height // 8, width // 8)
      UpperCamelCase       =  (batch_size * num_images_per_prompt, self.unet.config.in_channels, 6_4, 6_4)
      UpperCamelCase       =  text_embeddings.dtype
      if latents is None:
         if self.device.type == "mps":
            # randn does not exist on mps
            UpperCamelCase       =  torch.randn(
                lowerCamelCase__ ,   generator=lowerCamelCase__ ,   device='''cpu''' ,   dtype=lowerCamelCase__   ).to(self.device   )
            UpperCamelCase       =  torch.randn(lowerCamelCase__ ,   generator=lowerCamelCase__ ,   device='''cpu''' ,   dtype=lowerCamelCase__   ).to(
                self.device   )
         else:
            UpperCamelCase       =  torch.randn(
                lowerCamelCase__ ,   generator=lowerCamelCase__ ,   device=self.device ,   dtype=lowerCamelCase__   )
            UpperCamelCase       =  torch.randn(lowerCamelCase__ ,   generator=lowerCamelCase__ ,   device=self.device ,   dtype=lowerCamelCase__   )
      else:
         if latents_reference.shape != latents_shape:
            raise ValueError(f'Unexpected latents shape, got {latents.shape}, expected {latents_shape}'   )
         UpperCamelCase       =  latents_reference.to(self.device   )
         UpperCamelCase       =  latents.to(self.device   )

      # This is the key part of the pipeline where we
      # try to ensure that the generated images w/ the same seed
      # but different sizes actually result in similar images
      UpperCamelCase       =  (latents_shape[3] - latents_shape_reference[3]) // 2
      UpperCamelCase       =  (latents_shape[2] - latents_shape_reference[2]) // 2
      UpperCamelCase       =  latents_shape_reference[3] if dx >= 0 else latents_shape_reference[3] + 2 * dx
      UpperCamelCase       =  latents_shape_reference[2] if dy >= 0 else latents_shape_reference[2] + 2 * dy
      UpperCamelCase       =  0 if dx < 0 else dx
      UpperCamelCase       =  0 if dy < 0 else dy
      UpperCamelCase       =  max(-dx ,   0   )
      UpperCamelCase       =  max(-dy ,   0   )
      # import pdb
      # pdb.set_trace()
      UpperCamelCase       =  latents_reference[:, :, dy : dy + h, dx : dx + w]

      # set timesteps
      self.scheduler.set_timesteps(lowerCamelCase__   )

      # Some schedulers like PNDM have timesteps as arrays
      # It's more optimized to move all timesteps to correct device beforehand
      UpperCamelCase       =  self.scheduler.timesteps.to(self.device   )

      # scale the initial noise by the standard deviation required by the scheduler
      UpperCamelCase       =  latents * self.scheduler.init_noise_sigma

      # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature
      # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers.
      # eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502
      # and should be between [0, 1]
      UpperCamelCase       =  '''eta''' in set(inspect.signature(self.scheduler.step   ).parameters.keys()   )
      UpperCamelCase       =  {}
      if accepts_eta:
         UpperCamelCase       =  eta

      for i, t in enumerate(self.progress_bar(lowerCamelCase__   )   ):
         # expand the latents if we are doing classifier free guidance
         UpperCamelCase       =  torch.cat([latents] * 2   ) if do_classifier_free_guidance else latents
         UpperCamelCase       =  self.scheduler.scale_model_input(lowerCamelCase__ ,   lowerCamelCase__   )

         # predict the noise residual
         UpperCamelCase       =  self.unet(lowerCamelCase__ ,   lowerCamelCase__ ,   encoder_hidden_states=lowerCamelCase__   ).sample

         # perform guidance
         if do_classifier_free_guidance:
            UpperCamelCase     , UpperCamelCase       =  noise_pred.chunk(2   )
            UpperCamelCase       =  noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)

         # compute the previous noisy sample x_t -> x_t-1
         UpperCamelCase       =  self.scheduler.step(lowerCamelCase__ ,   lowerCamelCase__ ,   lowerCamelCase__ ,   **lowerCamelCase__   ).prev_sample

         # call the callback, if provided
         if callback is not None and i % callback_steps == 0:
            callback(lowerCamelCase__ ,   lowerCamelCase__ ,   lowerCamelCase__   )

      UpperCamelCase       =  1 / 0.18215 * latents
      UpperCamelCase       =  self.vae.decode(lowerCamelCase__   ).sample

      UpperCamelCase       =  (image / 2 + 0.5).clamp(0 ,   1   )

      # we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16
      UpperCamelCase       =  image.cpu().permute(0 ,   2 ,   3 ,   1   ).float().numpy()

      if self.safety_checker is not None:
         UpperCamelCase       =  self.feature_extractor(self.numpy_to_pil(lowerCamelCase__   ) ,   return_tensors='''pt'''   ).to(
             self.device   )
         UpperCamelCase     , UpperCamelCase       =  self.safety_checker(
             images=lowerCamelCase__ ,   clip_input=safety_checker_input.pixel_values.to(text_embeddings.dtype   )   )
      else:
         UpperCamelCase       =  None

      if output_type == "pil":
         UpperCamelCase       =  self.numpy_to_pil(lowerCamelCase__   )

      if not return_dict:
         return (image, has_nsfw_concept)

      return StableDiffusionPipelineOutput(images=lowerCamelCase__ ,   nsfw_content_detected=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()

snake_case_ :  Tuple     =					[
    '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_ :  Union[str, Any]     =					[
    'mlp.dense_4h_to_h.weight',
    'self_attention.dense.weight',
]



def  __snake_case     (    _UpperCAmelCase					:       List[str],      _UpperCAmelCase					:       List[str]):
   UpperCamelCase       =  {
       '''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
   UpperCamelCase       =  int(re.match(R'''.*layer_(\d*).*''',      _UpperCAmelCase)[1])
   layer_number -= 3
   return f'h.{layer_number}.' + key



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



def  __snake_case     (    _UpperCAmelCase					:       int,      _UpperCAmelCase					:       Optional[Any],      _UpperCAmelCase					:       str,      _UpperCAmelCase					:       List[Any],      _UpperCAmelCase					:       Optional[int]):
   # Construct model
   if bloom_config_file == "":
      UpperCamelCase       =  BloomConfig()
   else:
      UpperCamelCase       =  BloomConfig.from_json_file(_UpperCAmelCase)

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

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

      UpperCamelCase       =  None

      UpperCamelCase       =  BloomConfig()

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

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

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

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

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

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

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

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

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

         UpperCamelCase       =  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:
            UpperCamelCase       =  set(other_keys.missing_keys)
         else:
            UpperCamelCase       =  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)
      UpperCamelCase       =  pytorch_dump_folder_path + '''/''' + WEIGHTS_NAME
      UpperCamelCase       =  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:
         UpperCamelCase       =  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_ :  Tuple     =					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_ :  List[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,
   )





