






import gc
import unittest

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

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    StableDiffusionSAGPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

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


enable_full_determinism()


class     _a							(	UpperCAmelCase__    ,							UpperCAmelCase__    ,							unittest.TestCase     ):






  """simple docstring"""





  A_						       = StableDiffusionSAGPipeline
  A_						       = TEXT_TO_IMAGE_PARAMS
  A_						       = TEXT_TO_IMAGE_BATCH_PARAMS
  A_						       = TEXT_TO_IMAGE_IMAGE_PARAMS
  A_						       = TEXT_TO_IMAGE_IMAGE_PARAMS
  A_						       = False




  def  _UpperCAmelCase  (  self			)      ->			Optional[Any]:
     torch.manual_seed(0			)
     UpperCamelCase_						=  UNetaDConditionModel(
         block_out_channels=(32, 64)		,		layers_per_block=2		,		sample_size=32		,		in_channels=4		,		out_channels=4		,		down_block_types=('DownBlock2D', 'CrossAttnDownBlock2D')		,		up_block_types=('CrossAttnUpBlock2D', 'UpBlock2D')		,		cross_attention_dim=32		,		)
     UpperCamelCase_						=  DDIMScheduler(
         beta_start=0.0_0_0_8_5		,		beta_end=0.0_1_2		,		beta_schedule='scaled_linear'		,		clip_sample=_UpperCAmelCase		,		set_alpha_to_one=_UpperCAmelCase		,		)
     torch.manual_seed(0			)
     UpperCamelCase_						=  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		,		)
     torch.manual_seed(0			)
     UpperCamelCase_						=  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=1000		,		)
     UpperCamelCase_						=  CLIPTextModel(_UpperCAmelCase			)
     UpperCamelCase_						=  CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'			)

     UpperCamelCase_						=  {
         'unet': unet,
         'scheduler': scheduler,
         'vae': vae,
         'text_encoder': text_encoder,
         'tokenizer': tokenizer,
         'safety_checker': None,
         'feature_extractor': None,
     }
     return components




  def  _UpperCAmelCase  (  self		,		_UpperCAmelCase		,		_UpperCAmelCase=0			)      ->			List[Any]:
     if str(_UpperCAmelCase			).startswith('mps'			):
        UpperCamelCase_						=  torch.manual_seed(_UpperCAmelCase			)
     else:
        UpperCamelCase_						=  torch.Generator(device=_UpperCAmelCase			).manual_seed(_UpperCAmelCase			)
     UpperCamelCase_						=  {
         'prompt': '.',
         'generator': generator,
         'num_inference_steps': 2,
         'guidance_scale': 1.0,
         'sag_scale': 1.0,
         'output_type': 'numpy',
     }
     return inputs




  def  _UpperCAmelCase  (  self			)      ->			Tuple:
     super().test_inference_batch_single_identical(expected_max_diff=3e-3			)




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






  """simple docstring"""





  def  _UpperCAmelCase  (  self			)      ->			Tuple:
     # clean up the VRAM after each test
     super().tearDown()
     gc.collect()
     torch.cuda.empty_cache()




  def  _UpperCAmelCase  (  self			)      ->			str:
     UpperCamelCase_						=  StableDiffusionSAGPipeline.from_pretrained('CompVis/stable-diffusion-v1-4'			)
     UpperCamelCase_						=  sag_pipe.to(_UpperCAmelCase			)
     sag_pipe.set_progress_bar_config(disable=_UpperCAmelCase			)

     UpperCamelCase_						=  '.'
     UpperCamelCase_						=  torch.manual_seed(0			)
     UpperCamelCase_						=  sag_pipe(
         [prompt]		,		generator=_UpperCAmelCase		,		guidance_scale=7.5		,		sag_scale=1.0		,		num_inference_steps=20		,		output_type='np'			)

     UpperCamelCase_						=  output.images

     UpperCamelCase_						=  image[0, -3:, -3:, -1]

     assert image.shape == (1, 512, 512, 3)
     UpperCamelCase_						=  np.array([0.1_5_6_8, 0.1_7_3_8, 0.1_6_9_5, 0.1_6_9_3, 0.1_5_0_7, 0.1_7_0_5, 0.1_5_4_7, 0.1_7_5_1, 0.1_9_4_9]			)

     assert np.abs(image_slice.flatten() - expected_slice			).max() < 5e-2




  def  _UpperCAmelCase  (  self			)      ->			List[Any]:
     UpperCamelCase_						=  StableDiffusionSAGPipeline.from_pretrained('stabilityai/stable-diffusion-2-1-base'			)
     UpperCamelCase_						=  sag_pipe.to(_UpperCAmelCase			)
     sag_pipe.set_progress_bar_config(disable=_UpperCAmelCase			)

     UpperCamelCase_						=  '.'
     UpperCamelCase_						=  torch.manual_seed(0			)
     UpperCamelCase_						=  sag_pipe(
         [prompt]		,		generator=_UpperCAmelCase		,		guidance_scale=7.5		,		sag_scale=1.0		,		num_inference_steps=20		,		output_type='np'			)

     UpperCamelCase_						=  output.images

     UpperCamelCase_						=  image[0, -3:, -3:, -1]

     assert image.shape == (1, 512, 512, 3)
     UpperCamelCase_						=  np.array([0.3_4_5_9, 0.2_8_7_6, 0.2_5_3_7, 0.3_0_0_2, 0.2_6_7_1, 0.2_1_6_0, 0.3_0_2_6, 0.2_2_6_2, 0.2_3_7_1]			)

     assert np.abs(image_slice.flatten() - expected_slice			).max() < 5e-2




  def  _UpperCAmelCase  (  self			)      ->			Dict:
     UpperCamelCase_						=  StableDiffusionSAGPipeline.from_pretrained('stabilityai/stable-diffusion-2-1-base'			)
     UpperCamelCase_						=  sag_pipe.to(_UpperCAmelCase			)
     sag_pipe.set_progress_bar_config(disable=_UpperCAmelCase			)

     UpperCamelCase_						=  '.'
     UpperCamelCase_						=  torch.manual_seed(0			)
     UpperCamelCase_						=  sag_pipe(
         [prompt]		,		width=768		,		height=512		,		generator=_UpperCAmelCase		,		guidance_scale=7.5		,		sag_scale=1.0		,		num_inference_steps=20		,		output_type='np'		,		)

     UpperCamelCase_						=  output.images

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












'''simple docstring'''

import argparse
import re

from flax.traverse_util import flatten_dict, unflatten_dict
from tax import checkpoints

from transformers import SwitchTransformersConfig, SwitchTransformersForConditionalGeneration
from transformers.modeling_flax_pytorch_utils import load_flax_weights_in_pytorch_model
from transformers.utils import logging


logging.set_verbosity_info()


# should not include what is already done by the `from_pt` argument
a   	=    {
    "/attention/": "/0/SelfAttention/",
    "/self_attention/": "/0/SelfAttention/",
    "/encoder_decoder_attention/": "/1/EncDecAttention/",
    "value": "v",
    "query": "q",
    "key": "k",
    "out": "o",
    "pre_self_attention_layer_norm": "0/layer_norm",
    "pre_cross_attention_layer_norm": "1/layer_norm",
    "pre_attention_layer_norm": "0/layer_norm",  # previously 1, but seems wrong
    "token_embedder": "shared",
    "encoder_norm": "final_layer_norm",
    "decoder_norm": "final_layer_norm",
    "relpos_bias/rel_embedding": "block/0/layer/0/SelfAttention/relative_attention_bias/weight",
    "router/router_weights/w/": "router/classifier/",
    "roer/roer_weights/w/": "router/classifier/",
    "logits_dense": "lm_head",
}

def     __magic_name__							(				__UpperCAmelCase   )    ->       Optional[int]:

   '''simple docstring'''







   __SCREAMING_SNAKE_CASE   				=							list(s_dict.keys()   )
   for key in keys:
      __SCREAMING_SNAKE_CASE   				=							R""".*/layers_(\d+)"""
      __SCREAMING_SNAKE_CASE   				=							key
      if re.match(__UpperCAmelCase				,						__UpperCAmelCase   ):
         __SCREAMING_SNAKE_CASE   				=							re.sub(R"""layers_(\d+)"""				,						R"""block/\1/layer"""				,						__UpperCAmelCase   )

      __SCREAMING_SNAKE_CASE   				=							R"""(encoder|decoder)\/"""

      if re.match(__UpperCAmelCase				,						__UpperCAmelCase   ):
         __SCREAMING_SNAKE_CASE   				=							re.match(__UpperCAmelCase				,						__UpperCAmelCase   ).groups()
         if groups[0] == "encoder":
            __SCREAMING_SNAKE_CASE   				=							re.sub(R"""/mlp/"""				,						R"""/1/mlp/"""				,						__UpperCAmelCase   )
            __SCREAMING_SNAKE_CASE   				=							re.sub(R"""/pre_mlp_layer_norm/"""				,						R"""/1/layer_norm/"""				,						__UpperCAmelCase   )

         elif groups[0] == "decoder":
            __SCREAMING_SNAKE_CASE   				=							re.sub(R"""/mlp/"""				,						R"""/2/mlp/"""				,						__UpperCAmelCase   )
            __SCREAMING_SNAKE_CASE   				=							re.sub(R"""/pre_mlp_layer_norm/"""				,						R"""/2/layer_norm/"""				,						__UpperCAmelCase   )

        # 2. Convert other classic mappings
      for old_key, temp_key in MOE_LAYER_NAME_MAPPING.items():
         if old_key in new_key:
            __SCREAMING_SNAKE_CASE   				=							new_key.replace(__UpperCAmelCase				,						__UpperCAmelCase   )

      print(f"""{key} -> {new_key}"""   )
      __SCREAMING_SNAKE_CASE   				=							s_dict.pop(__UpperCAmelCase   )

   if "encoder/block/0/layer/0/SelfAttention/relative_attention_bias/weight" in s_dict:
      __SCREAMING_SNAKE_CASE   				=							s_dict[
          """encoder/block/0/layer/0/SelfAttention/relative_attention_bias/weight"""
      ].T
   if "decoder/block/0/layer/0/SelfAttention/relative_attention_bias/weight" in s_dict:
      __SCREAMING_SNAKE_CASE   				=							s_dict[
          """decoder/block/0/layer/0/SelfAttention/relative_attention_bias/weight"""
      ].T

   # 3. Take extra care of the EXPERTS layer
   for key in list(s_dict.keys()   ):
      if "expert" in key:
         __SCREAMING_SNAKE_CASE   				=							s_dict[key].shape[0]
         __SCREAMING_SNAKE_CASE   				=							s_dict[key]
         for idx in range(__UpperCAmelCase   ):
            __SCREAMING_SNAKE_CASE   				=							expert_weihts[idx]
            print(f"""{key} -> {key.replace('expert/'				,						'nested fstring'   )}"""   )

         s_dict.pop(__UpperCAmelCase   )

   return s_dict


a   	=    {
    "NUM_ENCODER_LAYERS": "num_layers",
    "NUM_DECODER_LAYERS": "num_decoder_layers",
    "NUM_HEADS": "num_heads",
    "HEAD_DIM": "d_kv",
    "EMBED_DIM": "d_model",
    "MLP_DIM": "d_ff",
    "NUM_SELECTED_EXPERTS": "num_selected_experts",
    "NUM_ENCODER_SPARSE_LAYERS": "num_sparse_encoder_layers",
    "NUM_DECODER_SPARSE_LAYERS": "num_sparse_decoder_layers",
    "dense.MlpBlock.activations": "feed_forward_proj",
}

def     __magic_name__							(				__UpperCAmelCase				,						__UpperCAmelCase   )    ->       str:

   '''simple docstring'''







   import regex as re

   with open(__UpperCAmelCase				,						"""r"""   ) as f:
      __SCREAMING_SNAKE_CASE   				=							f.read()

   __SCREAMING_SNAKE_CASE   				=							re.findall(R"""(.*) = ([0-9.]*)"""				,						__UpperCAmelCase   )
   __SCREAMING_SNAKE_CASE   				=							{}
   for param, value in regex_match:
      if param in GIN_TO_CONFIG_MAPPING and value != "":
         __SCREAMING_SNAKE_CASE   				=							float(__UpperCAmelCase   ) if """.""" in value else int(__UpperCAmelCase   )

   __SCREAMING_SNAKE_CASE   				=							re.findall(R"""(.*activations) = \(\'(.*)\',\)"""				,						__UpperCAmelCase   )[0]
   __SCREAMING_SNAKE_CASE   				=							str(activation[1]   )

   __SCREAMING_SNAKE_CASE   				=							num_experts
   __SCREAMING_SNAKE_CASE   				=							SwitchTransformersConfig(**__UpperCAmelCase   )
   return config

def     __magic_name__							(				__UpperCAmelCase				,						__UpperCAmelCase				,						__UpperCAmelCase=None				,						__UpperCAmelCase="./"				,						__UpperCAmelCase=8   )    ->       str:

   '''simple docstring'''







   print(f"""Loading flax weights from : {flax_checkpoint_path}"""   )
   __SCREAMING_SNAKE_CASE   				=							checkpoints.load_tax_checkpoint(__UpperCAmelCase   )

   if gin_file is not None:
      __SCREAMING_SNAKE_CASE   				=							convert_gin_to_config(__UpperCAmelCase				,						__UpperCAmelCase   )
   else:
      __SCREAMING_SNAKE_CASE   				=							SwitchTransformersConfig.from_pretrained(__UpperCAmelCase   )

   __SCREAMING_SNAKE_CASE   				=							SwitchTransformersForConditionalGeneration(__UpperCAmelCase   )

   __SCREAMING_SNAKE_CASE   				=							flax_params["""target"""]
   __SCREAMING_SNAKE_CASE   				=							flatten_dict(__UpperCAmelCase				,						sep="""/"""   )
   __SCREAMING_SNAKE_CASE   				=							rename_keys(__UpperCAmelCase   )
   __SCREAMING_SNAKE_CASE   				=							unflatten_dict(__UpperCAmelCase				,						sep="""/"""   )

   # Load the flax params in the PT model
   load_flax_weights_in_pytorch_model(__UpperCAmelCase				,						__UpperCAmelCase   )

   print(f"""Save PyTorch model to {pytorch_dump_path}"""   )
   pt_model.save_pretrained(__UpperCAmelCase   )


if __name__ == "__main__":
    a   	=    argparse.ArgumentParser()
    # Required parameters
    parser.add_argument(
        "--switch_t5x_checkpoint_path",
        default=None,
        type=str,
        required=True,
        help=(
            "The config json file corresponding to the pre-trained SwitchTransformers model. \nThis specifies the"
            " model architecture. If not provided, a `gin_file` has to be provided."
        ),
    )
    parser.add_argument(
        "--gin_file",
        default=None,
        type=str,
        required=False,
        help="Path to the gin config file. If not provided, a `config_file` has to be passed   ",
    )
    parser.add_argument(
        "--config_name", default=None, type=str, required=False, help="Config name of SwitchTransformers model."
    )
    parser.add_argument(
        "--pytorch_dump_folder_path", default=None, type=str, required=True, help="Path to the output pytorch model."
    )
    parser.add_argument("--num_experts", default=8, type=int, required=False, help="Number of experts")
    a   	=    parser.parse_args()
    convert_flax_checkpoint_to_pytorch(
        args.switch_tax_checkpoint_path,
        args.config_name,
        args.gin_file,
        args.pytorch_dump_folder_path,
        args.num_experts,
    )





