



"""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
lowerCAmelCase__      =    {
    '''/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 	snake_case_  (				A_ :		str						):


       '''simple docstring'''





       _lowerCamelCase						:    Any         =			list(s_dict.keys()						)
       for key in keys:
              _lowerCamelCase						:    Dict         =			R'''.*/layers_(\d+)'''
              _lowerCamelCase						:    Union[str, Any]         =			key
              if re.match(A_,    A_						):
                     _lowerCamelCase						:    str         =			re.sub(R'''layers_(\d+)''',    R'''block/\1/layer''',    A_						)

              _lowerCamelCase						:    int         =			R'''(encoder|decoder)\/'''

              if re.match(A_,    A_						):
                     _lowerCamelCase						:    Tuple         =			re.match(A_,    A_						).groups()
                     if groups[0] == "encoder":
                            _lowerCamelCase						:    Optional[Any]         =			re.sub(R'''/mlp/''',    R'''/1/mlp/''',    A_						)
                            _lowerCamelCase						:    Optional[Any]         =			re.sub(R'''/pre_mlp_layer_norm/''',    R'''/1/layer_norm/''',    A_						)

                     elif groups[0] == "decoder":
                            _lowerCamelCase						:    Any         =			re.sub(R'''/mlp/''',    R'''/2/mlp/''',    A_						)
                            _lowerCamelCase						:    Optional[int]         =			re.sub(R'''/pre_mlp_layer_norm/''',    R'''/2/layer_norm/''',    A_						)

        # 2. Convert other classic mappings
              for old_key, temp_key in MOE_LAYER_NAME_MAPPING.items():
                     if old_key in new_key:
                            _lowerCamelCase						:    List[str]         =			new_key.replace(A_,    A_						)

              print(F'''{key} -> {new_key}'''						)
              _lowerCamelCase						:    int         =			s_dict.pop(A_						)

       if "encoder/block/0/layer/0/SelfAttention/relative_attention_bias/weight" in s_dict:
              _lowerCamelCase						:    List[Any]         =			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:
              _lowerCamelCase						:    List[str]         =			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:
                     _lowerCamelCase						:    Optional[Any]         =			s_dict[key].shape[0]
                     _lowerCamelCase						:    List[str]         =			s_dict[key]
                     for idx in range(A_						):
                            _lowerCamelCase						:    Optional[Any]         =			expert_weihts[idx]
                            print(F'''{key} -> {key.replace("expert/",    "nested fstring"						)}'''						)

                     s_dict.pop(A_						)

       return s_dict


lowerCAmelCase__      =    {
    '''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 	snake_case_  (				A_ :		List[str],    A_ :		str						):


       '''simple docstring'''





       import regex as re

       with open(A_,    '''r'''						) as f:
              _lowerCamelCase						:    Dict         =			f.read()

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

       _lowerCamelCase						:    Any         =			re.findall(R'''(.*activations) = \(\'(.*)\',\)''',    A_						)[0]
       _lowerCamelCase						:    Tuple         =			str(activation[1]						)

       _lowerCamelCase						:    Any         =			num_experts
       _lowerCamelCase						:    Optional[int]         =			SwitchTransformersConfig(**A_						)
       return config






def 	snake_case_  (				A_ :		Tuple,    A_ :		Tuple,    A_ :		Tuple=None,    A_ :		Any="./",    A_ :		Optional[int]=8						):


       '''simple docstring'''





       print(F'''Loading flax weights from : {flax_checkpoint_path}'''						)
       _lowerCamelCase						:    Dict         =			checkpoints.load_tax_checkpoint(A_						)

       if gin_file is not None:
              _lowerCamelCase						:    Dict         =			convert_gin_to_config(A_,    A_						)
       else:
              _lowerCamelCase						:    List[str]         =			SwitchTransformersConfig.from_pretrained(A_						)

       _lowerCamelCase						:    int         =			SwitchTransformersForConditionalGeneration(A_						)

       _lowerCamelCase						:    List[str]         =			flax_params['''target''']
       _lowerCamelCase						:    Any         =			flatten_dict(A_,    sep='''/'''						)
       _lowerCamelCase						:    Optional[int]         =			rename_keys(A_						)
       _lowerCamelCase						:    Dict         =			unflatten_dict(A_,    sep='''/'''						)

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

       print(F'''Save PyTorch model to {pytorch_dump_path}'''						)
       pt_model.save_pretrained(A_						)


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







"""simple docstring"""






import random
import unittest

import numpy as np
import torch

from diffusers import (
    DPMSolverMultistepScheduler,
    EulerAncestralDiscreteScheduler,
    EulerDiscreteScheduler,
    LMSDiscreteScheduler,
    OnnxStableDiffusionUpscalePipeline,
    PNDMScheduler,
)
from diffusers.utils import floats_tensor
from diffusers.utils.testing_utils import (
    is_onnx_available,
    load_image,
    nightly,
    require_onnxruntime,
    require_torch_gpu,
)

from ..test_pipelines_onnx_common import OnnxPipelineTesterMixin


if is_onnx_available():
   import onnxruntime as ort






class 	__snake_case	(     _lowercase  ,					unittest.TestCase):
  # TODO: is there an appropriate internal test set?
  snake_case__    :       List[str]        =						"ssube/stable-diffusion-x4-upscaler-onnx"




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






         """simple docstring"""






         _lowerCamelCase						:    Tuple         =			floats_tensor((1, 3, 1_2_8, 1_2_8)						, rng=random.Random(__lowerCAmelCase		)		)
         _lowerCamelCase						:    Union[str, Any]         =			torch.manual_seed(__lowerCAmelCase		)
         _lowerCamelCase						:    Tuple         =			{
             '''prompt''': '''A painting of a squirrel eating a burger''',
             '''image''': image,
             '''generator''': generator,
             '''num_inference_steps''': 3,
             '''guidance_scale''': 7.5,
             '''output_type''': '''numpy''',
         }
         return inputs




  def 						SCREAMING_SNAKE_CASE    (			self					:						Optional[Any]		):






         """simple docstring"""






         _lowerCamelCase						:    Union[str, Any]         =			OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint						, provider='''CPUExecutionProvider'''		)
         pipe.set_progress_bar_config(disable=__lowerCAmelCase		)

         _lowerCamelCase						:    Union[str, Any]         =			self.get_dummy_inputs()
         _lowerCamelCase						:    Any         =			pipe(**__lowerCAmelCase		).images
         _lowerCamelCase						:    Dict         =			image[0, -3:, -3:, -1].flatten()

         # started as 128, should now be 512
         assert image.shape == (1, 5_1_2, 5_1_2, 3)
         _lowerCamelCase						:    str         =			np.array(
             [0.6_97_47_82, 0.68_90_20_93, 0.70_13_58_85, 0.7_58_36_18, 0.7_80_45_45, 0.7_85_49_12, 0.78_66_74_26, 0.78_74_38_63, 0.78_07_02_23]		)
         assert np.abs(image_slice - expected_slice		).max() < 1E-1




  def 						SCREAMING_SNAKE_CASE    (			self					:						List[Any]		):






         """simple docstring"""






         _lowerCamelCase						:    Optional[Any]         =			OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint						, provider='''CPUExecutionProvider'''		)
         _lowerCamelCase						:    List[str]         =			PNDMScheduler.from_config(pipe.scheduler.config						, skip_prk_steps=__lowerCAmelCase		)
         pipe.set_progress_bar_config(disable=__lowerCAmelCase		)

         _lowerCamelCase						:    int         =			self.get_dummy_inputs()
         _lowerCamelCase						:    Optional[Any]         =			pipe(**__lowerCAmelCase		).images
         _lowerCamelCase						:    str         =			image[0, -3:, -3:, -1]

         assert image.shape == (1, 5_1_2, 5_1_2, 3)
         _lowerCamelCase						:    Optional[int]         =			np.array(
             [0.6_89_88_92, 0.59_24_05_56, 0.52_49_95_27, 0.58_86_62_15, 0.52_25_82_35, 0.52_57_27_15, 0.62_41_44_73, 0.6_17_43_87, 0.6_21_49_64]		)
         assert np.abs(image_slice.flatten() - expected_slice		).max() < 1E-1




  def 						SCREAMING_SNAKE_CASE    (			self					:						Optional[int]		):






         """simple docstring"""






         _lowerCamelCase						:    Optional[int]         =			OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint						, provider='''CPUExecutionProvider'''		)
         _lowerCamelCase						:    Optional[int]         =			DPMSolverMultistepScheduler.from_config(pipe.scheduler.config		)
         pipe.set_progress_bar_config(disable=__lowerCAmelCase		)

         _lowerCamelCase						:    Tuple         =			self.get_dummy_inputs()
         _lowerCamelCase						:    str         =			pipe(**__lowerCAmelCase		).images
         _lowerCamelCase						:    Tuple         =			image[0, -3:, -3:, -1]

         assert image.shape == (1, 5_1_2, 5_1_2, 3)
         _lowerCamelCase						:    str         =			np.array(
             [0.7_65_92_78, 0.76_43_76_64, 0.75_57_91_07, 0.7_69_11_16, 0.77_66_69_86, 0.7_72_76_72, 0.7_75_86_64, 0.7_81_22_26, 0.76_94_25_15]		)

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




  def 						SCREAMING_SNAKE_CASE    (			self					:						Any		):






         """simple docstring"""






         _lowerCamelCase						:    Any         =			OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint						, provider='''CPUExecutionProvider'''		)
         _lowerCamelCase						:    Union[str, Any]         =			EulerDiscreteScheduler.from_config(pipe.scheduler.config		)
         pipe.set_progress_bar_config(disable=__lowerCAmelCase		)

         _lowerCamelCase						:    Optional[Any]         =			self.get_dummy_inputs()
         _lowerCamelCase						:    Tuple         =			pipe(**__lowerCAmelCase		).images
         _lowerCamelCase						:    Union[str, Any]         =			image[0, -3:, -3:, -1]

         assert image.shape == (1, 5_1_2, 5_1_2, 3)
         _lowerCamelCase						:    Union[str, Any]         =			np.array(
             [0.6_97_47_82, 0.68_90_20_93, 0.70_13_58_85, 0.7_58_36_18, 0.7_80_45_45, 0.7_85_49_12, 0.78_66_74_26, 0.78_74_38_63, 0.78_07_02_23]		)
         assert np.abs(image_slice.flatten() - expected_slice		).max() < 1E-1




  def 						SCREAMING_SNAKE_CASE    (			self					:						Any		):






         """simple docstring"""






         _lowerCamelCase						:    Dict         =			OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint						, provider='''CPUExecutionProvider'''		)
         _lowerCamelCase						:    int         =			EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config		)
         pipe.set_progress_bar_config(disable=__lowerCAmelCase		)

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

         assert image.shape == (1, 5_1_2, 5_1_2, 3)
         _lowerCamelCase						:    Optional[int]         =			np.array(
             [0.77_42_44_96, 0.77_36_01, 0.7_64_52_88, 0.7_76_95_98, 0.7_77_27_39, 0.7_73_86_88, 0.78_18_72_33, 0.77_87_95_84, 0.76_70_43]		)

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






@nightly
@require_onnxruntime
@require_torch_gpu
class 	__snake_case	(     unittest.TestCase):




  @property
  def 						SCREAMING_SNAKE_CASE    (			self					:						str		):






         """simple docstring"""






         return (
             "CUDAExecutionProvider",
             {
                 "gpu_mem_limit": "15000000000",  # 15GB
                 "arena_extend_strategy": "kSameAsRequested",
             },
         )




  @property
  def 						SCREAMING_SNAKE_CASE    (			self					:						Any		):






         """simple docstring"""






         _lowerCamelCase						:    Optional[int]         =			ort.SessionOptions()
         _lowerCamelCase						:    List[str]         =			False
         return options




  def 						SCREAMING_SNAKE_CASE    (			self					:						Optional[int]		):






         """simple docstring"""






         _lowerCamelCase						:    Optional[int]         =			load_image(
             '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
             '''/img2img/sketch-mountains-input.jpg'''		)
         _lowerCamelCase						:    Any         =			init_image.resize((1_2_8, 1_2_8)		)
         # using the PNDM scheduler by default
         _lowerCamelCase						:    List[str]         =			OnnxStableDiffusionUpscalePipeline.from_pretrained(
             '''ssube/stable-diffusion-x4-upscaler-onnx'''						, provider=self.gpu_provider						, sess_options=self.gpu_options						, )
         pipe.set_progress_bar_config(disable=__lowerCAmelCase		)

         _lowerCamelCase						:    int         =			'''A fantasy landscape, trending on artstation'''

         _lowerCamelCase						:    List[Any]         =			torch.manual_seed(0		)
         _lowerCamelCase						:    List[str]         =			pipe(
             prompt=__lowerCAmelCase						, image=__lowerCAmelCase						, guidance_scale=7.5						, num_inference_steps=1_0						, generator=__lowerCAmelCase						, output_type='''np'''						, )
         _lowerCamelCase						:    List[Any]         =			output.images
         _lowerCamelCase						:    List[Any]         =			images[0, 2_5_5:2_5_8, 3_8_3:3_8_6, -1]

         assert images.shape == (1, 5_1_2, 5_1_2, 3)
         _lowerCamelCase						:    str         =			np.array([0.48_83, 0.49_47, 0.49_80, 0.49_75, 0.49_82, 0.49_80, 0.50_00, 0.50_06, 0.49_72]		)
         # TODO: lower the tolerance after finding the cause of onnxruntime reproducibility issues

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




  def 						SCREAMING_SNAKE_CASE    (			self					:						str		):






         """simple docstring"""






         _lowerCamelCase						:    Union[str, Any]         =			load_image(
             '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
             '''/img2img/sketch-mountains-input.jpg'''		)
         _lowerCamelCase						:    int         =			init_image.resize((1_2_8, 1_2_8)		)
         _lowerCamelCase						:    str         =			LMSDiscreteScheduler.from_pretrained(
             '''ssube/stable-diffusion-x4-upscaler-onnx'''						, subfolder='''scheduler'''		)
         _lowerCamelCase						:    Dict         =			OnnxStableDiffusionUpscalePipeline.from_pretrained(
             '''ssube/stable-diffusion-x4-upscaler-onnx'''						, scheduler=__lowerCAmelCase						, provider=self.gpu_provider						, sess_options=self.gpu_options						, )
         pipe.set_progress_bar_config(disable=__lowerCAmelCase		)

         _lowerCamelCase						:    Optional[int]         =			'''A fantasy landscape, trending on artstation'''

         _lowerCamelCase						:    int         =			torch.manual_seed(0		)
         _lowerCamelCase						:    List[str]         =			pipe(
             prompt=__lowerCAmelCase						, image=__lowerCAmelCase						, guidance_scale=7.5						, num_inference_steps=2_0						, generator=__lowerCAmelCase						, output_type='''np'''						, )
         _lowerCamelCase						:    Union[str, Any]         =			output.images
         _lowerCamelCase						:    Optional[Any]         =			images[0, 2_5_5:2_5_8, 3_8_3:3_8_6, -1]

         assert images.shape == (1, 5_1_2, 5_1_2, 3)
         _lowerCamelCase						:    str         =			np.array(
             [0.50_17_37_53, 0.50_22_33_56, 0.50_20_39, 0.50_23_30_36, 0.5_02_37_25, 0.5_02_26_01, 0.5_01_87_58, 0.50_23_40_85, 0.50_24_15_66]		)
         # TODO: lower the tolerance after finding the cause of onnxruntime reproducibility issues

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


