


import gc
import random
import unittest

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

from diffusers import AutoencoderKL, DDIMScheduler, DDPMScheduler, StableDiffusionUpscalePipeline, UNetaDConditionModel
from diffusers.utils import floats_tensor, load_image, load_numpy, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu


enable_full_determinism()







class lowerCAmelCase_		(				unittest.TestCase ):

 '''simple docstring'''

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



 @property
 def 		UpperCamelCase__  (			self      ):
       snake_case_ =   1
       snake_case_ =   3
       snake_case_ =   (32, 32)

       snake_case_ =   floats_tensor((batch_size, num_channels) + sizes	,   rng=random.Random(0      )      ).to(_UpperCAmelCase      )
       return image



 @property
 def 		UpperCamelCase__  (			self      ):
       torch.manual_seed(0      )
       snake_case_ =   UNetaDConditionModel(
           block_out_channels=(32, 32, 64)	,   layers_per_block=2	,   sample_size=32	,   in_channels=7	,   out_channels=4	,   down_block_types=('''DownBlock2D''', '''CrossAttnDownBlock2D''', '''CrossAttnDownBlock2D''')	,   up_block_types=('''CrossAttnUpBlock2D''', '''CrossAttnUpBlock2D''', '''UpBlock2D''')	,   cross_attention_dim=32	,   attention_head_dim=8	,   use_linear_projection=_UpperCAmelCase	,   only_cross_attention=(True, True, False)	,   num_class_embeds=1_00	,   )
       return model



 @property
 def 		UpperCamelCase__  (			self      ):
       torch.manual_seed(0      )
       snake_case_ =   AutoencoderKL(
           block_out_channels=[32, 32, 64]	,   in_channels=3	,   out_channels=3	,   down_block_types=['''DownEncoderBlock2D''', '''DownEncoderBlock2D''', '''DownEncoderBlock2D''']	,   up_block_types=['''UpDecoderBlock2D''', '''UpDecoderBlock2D''', '''UpDecoderBlock2D''']	,   latent_channels=4	,   )
       return model



 @property
 def 		UpperCamelCase__  (			self      ):
       torch.manual_seed(0      )
       snake_case_ =   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=10_00	,   hidden_act='''gelu'''	,   projection_dim=5_12	,   )
       return CLIPTextModel(_UpperCAmelCase      )



 def 		UpperCamelCase__  (			self      ):
       snake_case_ =   '''cpu'''  # ensure determinism for the device-dependent torch.Generator
       snake_case_ =   self.dummy_cond_unet_upscale
       snake_case_ =   DDPMScheduler()
       snake_case_ =   DDIMScheduler(prediction_type='''v_prediction'''      )
       snake_case_ =   self.dummy_vae
       snake_case_ =   self.dummy_text_encoder
       snake_case_ =   CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''      )

       snake_case_ =   self.dummy_image.cpu().permute(0	,   2	,   3	,   1      )[0]
       snake_case_ =   Image.fromarray(np.uinta(_UpperCAmelCase      )      ).convert('''RGB'''      ).resize((64, 64)      )

       # make sure here that pndm scheduler skips prk
       snake_case_ =   StableDiffusionUpscalePipeline(
           unet=_UpperCAmelCase	,   low_res_scheduler=_UpperCAmelCase	,   scheduler=_UpperCAmelCase	,   vae=_UpperCAmelCase	,   text_encoder=_UpperCAmelCase	,   tokenizer=_UpperCAmelCase	,   max_noise_level=3_50	,   )
       snake_case_ =   sd_pipe.to(_UpperCAmelCase      )
       sd_pipe.set_progress_bar_config(disable=_UpperCAmelCase      )

       snake_case_ =   '''A painting of a squirrel eating a burger'''
       snake_case_ =   torch.Generator(device=_UpperCAmelCase      ).manual_seed(0      )
       snake_case_ =   sd_pipe(
           [prompt]	,   image=_UpperCAmelCase	,   generator=_UpperCAmelCase	,   guidance_scale=6.0	,   noise_level=20	,   num_inference_steps=2	,   output_type='''np'''	,   )

       snake_case_ =   output.images

       snake_case_ =   torch.Generator(device=_UpperCAmelCase      ).manual_seed(0      )
       snake_case_ =   sd_pipe(
           [prompt]	,   image=_UpperCAmelCase	,   generator=_UpperCAmelCase	,   guidance_scale=6.0	,   noise_level=20	,   num_inference_steps=2	,   output_type='''np'''	,   return_dict=_UpperCAmelCase	,   )[0]

       snake_case_ =   image[0, -3:, -3:, -1]
       snake_case_ =   image_from_tuple[0, -3:, -3:, -1]

       snake_case_ =   low_res_image.size[0] * 4
       assert image.shape == (1, expected_height_width, expected_height_width, 3)
       snake_case_ =   np.array([0.3_113, 0.3_910, 0.4_272, 0.4_859, 0.5_061, 0.4_652, 0.5_362, 0.5_715, 0.5_661]      )

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



 def 		UpperCamelCase__  (			self      ):
       snake_case_ =   '''cpu'''  # ensure determinism for the device-dependent torch.Generator
       snake_case_ =   self.dummy_cond_unet_upscale
       snake_case_ =   DDPMScheduler()
       snake_case_ =   DDIMScheduler(prediction_type='''v_prediction'''      )
       snake_case_ =   self.dummy_vae
       snake_case_ =   self.dummy_text_encoder
       snake_case_ =   CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''      )

       snake_case_ =   self.dummy_image.cpu().permute(0	,   2	,   3	,   1      )[0]
       snake_case_ =   Image.fromarray(np.uinta(_UpperCAmelCase      )      ).convert('''RGB'''      ).resize((64, 64)      )

       # make sure here that pndm scheduler skips prk
       snake_case_ =   StableDiffusionUpscalePipeline(
           unet=_UpperCAmelCase	,   low_res_scheduler=_UpperCAmelCase	,   scheduler=_UpperCAmelCase	,   vae=_UpperCAmelCase	,   text_encoder=_UpperCAmelCase	,   tokenizer=_UpperCAmelCase	,   max_noise_level=3_50	,   )
       snake_case_ =   sd_pipe.to(_UpperCAmelCase      )
       sd_pipe.set_progress_bar_config(disable=_UpperCAmelCase      )

       snake_case_ =   '''A painting of a squirrel eating a burger'''
       snake_case_ =   sd_pipe(
           2 * [prompt]	,   image=2 * [low_res_image]	,   guidance_scale=6.0	,   noise_level=20	,   num_inference_steps=2	,   output_type='''np'''	,   )
       snake_case_ =   output.images
       assert image.shape[0] == 2

       snake_case_ =   torch.Generator(device=_UpperCAmelCase      ).manual_seed(0      )
       snake_case_ =   sd_pipe(
           [prompt]	,   image=_UpperCAmelCase	,   generator=_UpperCAmelCase	,   num_images_per_prompt=2	,   guidance_scale=6.0	,   noise_level=20	,   num_inference_steps=2	,   output_type='''np'''	,   )
       snake_case_ =   output.images
       assert image.shape[0] == 2





 @unittest.skipIf(torch_device != '''cuda'''	,   '''This test requires a GPU'''      )
 def 		UpperCamelCase__  (			self      ):
       snake_case_ =   self.dummy_cond_unet_upscale
       snake_case_ =   DDPMScheduler()
       snake_case_ =   DDIMScheduler(prediction_type='''v_prediction'''      )
       snake_case_ =   self.dummy_vae
       snake_case_ =   self.dummy_text_encoder
       snake_case_ =   CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''      )

       snake_case_ =   self.dummy_image.cpu().permute(0	,   2	,   3	,   1      )[0]
       snake_case_ =   Image.fromarray(np.uinta(_UpperCAmelCase      )      ).convert('''RGB'''      ).resize((64, 64)      )

       # put models in fp16, except vae as it overflows in fp16
       snake_case_ =   unet.half()
       snake_case_ =   text_encoder.half()

       # make sure here that pndm scheduler skips prk
       snake_case_ =   StableDiffusionUpscalePipeline(
           unet=_UpperCAmelCase	,   low_res_scheduler=_UpperCAmelCase	,   scheduler=_UpperCAmelCase	,   vae=_UpperCAmelCase	,   text_encoder=_UpperCAmelCase	,   tokenizer=_UpperCAmelCase	,   max_noise_level=3_50	,   )
       snake_case_ =   sd_pipe.to(_UpperCAmelCase      )
       sd_pipe.set_progress_bar_config(disable=_UpperCAmelCase      )

       snake_case_ =   '''A painting of a squirrel eating a burger'''
       snake_case_ =   torch.manual_seed(0      )
       snake_case_ =   sd_pipe(
           [prompt]	,   image=_UpperCAmelCase	,   generator=_UpperCAmelCase	,   num_inference_steps=2	,   output_type='''np'''	,   ).images

       snake_case_ =   low_res_image.size[0] * 4
       assert image.shape == (1, expected_height_width, expected_height_width, 3)







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

 '''simple docstring'''

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



 def 		UpperCamelCase__  (			self      ):
       snake_case_ =   load_image(
           '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
           '''/sd2-upscale/low_res_cat.png'''      )
       snake_case_ =   load_numpy(
           '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd2-upscale'''
           '''/upsampled_cat.npy'''      )

       snake_case_ =   '''stabilityai/stable-diffusion-x4-upscaler'''
       snake_case_ =   StableDiffusionUpscalePipeline.from_pretrained(_UpperCAmelCase      )
       pipe.to(_UpperCAmelCase      )
       pipe.set_progress_bar_config(disable=_UpperCAmelCase      )
       pipe.enable_attention_slicing()

       snake_case_ =   '''a cat sitting on a park bench'''

       snake_case_ =   torch.manual_seed(0      )
       snake_case_ =   pipe(
           prompt=_UpperCAmelCase	,   image=_UpperCAmelCase	,   generator=_UpperCAmelCase	,   output_type='''np'''	,   )
       snake_case_ =   output.images[0]

       assert image.shape == (5_12, 5_12, 3)
       assert np.abs(expected_image - image      ).max() < 1E-3



 def 		UpperCamelCase__  (			self      ):
       snake_case_ =   load_image(
           '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
           '''/sd2-upscale/low_res_cat.png'''      )
       snake_case_ =   load_numpy(
           '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd2-upscale'''
           '''/upsampled_cat_fp16.npy'''      )

       snake_case_ =   '''stabilityai/stable-diffusion-x4-upscaler'''
       snake_case_ =   StableDiffusionUpscalePipeline.from_pretrained(
           _UpperCAmelCase	,   torch_dtype=torch.floataa	,   )
       pipe.to(_UpperCAmelCase      )
       pipe.set_progress_bar_config(disable=_UpperCAmelCase      )
       pipe.enable_attention_slicing()

       snake_case_ =   '''a cat sitting on a park bench'''

       snake_case_ =   torch.manual_seed(0      )
       snake_case_ =   pipe(
           prompt=_UpperCAmelCase	,   image=_UpperCAmelCase	,   generator=_UpperCAmelCase	,   output_type='''np'''	,   )
       snake_case_ =   output.images[0]

       assert image.shape == (5_12, 5_12, 3)
       assert np.abs(expected_image - image      ).max() < 5E-1





 def 		UpperCamelCase__  (			self      ):
       torch.cuda.empty_cache()
       torch.cuda.reset_max_memory_allocated()
       torch.cuda.reset_peak_memory_stats()

       snake_case_ =   load_image(
           '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
           '''/sd2-upscale/low_res_cat.png'''      )

       snake_case_ =   '''stabilityai/stable-diffusion-x4-upscaler'''
       snake_case_ =   StableDiffusionUpscalePipeline.from_pretrained(
           _UpperCAmelCase	,   torch_dtype=torch.floataa	,   )
       pipe.to(_UpperCAmelCase      )
       pipe.set_progress_bar_config(disable=_UpperCAmelCase      )
       pipe.enable_attention_slicing(1      )
       pipe.enable_sequential_cpu_offload()

       snake_case_ =   '''a cat sitting on a park bench'''

       snake_case_ =   torch.manual_seed(0      )
       snake_case_ =   pipe(
           prompt=_UpperCAmelCase	,   image=_UpperCAmelCase	,   generator=_UpperCAmelCase	,   num_inference_steps=5	,   output_type='''np'''	,   )

       snake_case_ =   torch.cuda.max_memory_allocated()
       # make sure that less than 2.9 GB is allocated
       assert mem_bytes < 2.9 * 10**9



import argparse
import pickle

import numpy as np
import torch
from torch import nn

from transformers import ReformerConfig, ReformerModelWithLMHead
from transformers.utils import logging


logging.set_verbosity_info()






def  __lowerCAmelCase					(SCREAMING_SNAKE_CASE      ,      SCREAMING_SNAKE_CASE      ,      SCREAMING_SNAKE_CASE=None      )->						Any:




      """simple docstring"""

      assert torch_layer.weight.shape == weight.shape, f'''{torch_layer} layer.weight does not match'''
      snake_case_ =   nn.Parameter(SCREAMING_SNAKE_CASE      )
      if bias is not None:
            assert torch_layer.bias.shape == bias.shape, f'''{torch_layer} layer.bias does not match'''
            snake_case_ =   nn.Parameter(SCREAMING_SNAKE_CASE      )






def  __lowerCAmelCase					(SCREAMING_SNAKE_CASE      ,      SCREAMING_SNAKE_CASE      ,      SCREAMING_SNAKE_CASE      )->						int:




      """simple docstring"""

      snake_case_ =   np.asarray(weights[0]      )
      snake_case_ =   np.asarray(weights[1]      )
      snake_case_ =   np.asarray(weights[2]      )

      set_param(
          torch_layer.self_attention.query_key      ,      torch.tensor(SCREAMING_SNAKE_CASE      ).transpose(1      ,      2      ).contiguous().view(-1      ,      SCREAMING_SNAKE_CASE      )      ,      )
      set_param(
          torch_layer.self_attention.value      ,      torch.tensor(SCREAMING_SNAKE_CASE      ).transpose(1      ,      2      ).contiguous().view(-1      ,      SCREAMING_SNAKE_CASE      )      ,      )
      set_param(
          torch_layer.output.dense      ,      torch.tensor(SCREAMING_SNAKE_CASE      ).view(-1      ,      SCREAMING_SNAKE_CASE      ).contiguous().transpose(0      ,      1      )      ,      )






def  __lowerCAmelCase					(SCREAMING_SNAKE_CASE      ,      SCREAMING_SNAKE_CASE      ,      SCREAMING_SNAKE_CASE      )->						Tuple:




      """simple docstring"""

      snake_case_ =   np.asarray(weights[0]      )
      snake_case_ =   np.asarray(weights[1]      )
      snake_case_ =   np.asarray(weights[2]      )
      snake_case_ =   np.asarray(weights[3]      )

      set_param(
          torch_layer.self_attention.query      ,      torch.tensor(SCREAMING_SNAKE_CASE      ).transpose(1      ,      2      ).contiguous().view(-1      ,      SCREAMING_SNAKE_CASE      )      ,      )
      set_param(
          torch_layer.self_attention.key      ,      torch.tensor(SCREAMING_SNAKE_CASE      ).transpose(1      ,      2      ).contiguous().view(-1      ,      SCREAMING_SNAKE_CASE      )      ,      )
      set_param(
          torch_layer.self_attention.value      ,      torch.tensor(SCREAMING_SNAKE_CASE      ).transpose(1      ,      2      ).contiguous().view(-1      ,      SCREAMING_SNAKE_CASE      )      ,      )
      set_param(
          torch_layer.output.dense      ,      torch.tensor(SCREAMING_SNAKE_CASE      ).view(-1      ,      SCREAMING_SNAKE_CASE      ).contiguous().transpose(0      ,      1      )      ,      )






def  __lowerCAmelCase					(SCREAMING_SNAKE_CASE      ,      SCREAMING_SNAKE_CASE      ,      SCREAMING_SNAKE_CASE      )->						int:




      """simple docstring"""

      snake_case_ =   weights[0][0][0]
      snake_case_ =   np.asarray(layer_norm_a[0]      )
      snake_case_ =   np.asarray(layer_norm_a[1]      )
      set_param(
          torch_block.attention.layer_norm      ,      torch.tensor(SCREAMING_SNAKE_CASE      )      ,      torch.tensor(SCREAMING_SNAKE_CASE      )      ,      )

      # lsh weights + output
      snake_case_ =   weights[0][1]
      if len(SCREAMING_SNAKE_CASE      ) < 4:
            set_layer_weights_in_torch_lsh(SCREAMING_SNAKE_CASE      ,      torch_block.attention      ,      SCREAMING_SNAKE_CASE      )
      else:
            set_layer_weights_in_torch_local(SCREAMING_SNAKE_CASE      ,      torch_block.attention      ,      SCREAMING_SNAKE_CASE      )

      # intermediate weighs
      snake_case_ =   weights[2][0][1][2]

      # Chunked Feed Forward
      if len(SCREAMING_SNAKE_CASE      ) == 4:
            snake_case_ =   intermediate_weights[2]

      # layernorm 2
      snake_case_ =   np.asarray(intermediate_weights[0][0]      )
      snake_case_ =   np.asarray(intermediate_weights[0][1]      )
      set_param(
          torch_block.feed_forward.layer_norm      ,      torch.tensor(SCREAMING_SNAKE_CASE      )      ,      torch.tensor(SCREAMING_SNAKE_CASE      )      ,      )

      # intermediate dense
      snake_case_ =   np.asarray(intermediate_weights[1][0]      )
      snake_case_ =   np.asarray(intermediate_weights[1][1]      )
      set_param(
          torch_block.feed_forward.dense.dense      ,      torch.tensor(SCREAMING_SNAKE_CASE      ).transpose(0      ,      1      ).contiguous()      ,      torch.tensor(SCREAMING_SNAKE_CASE      )      ,      )

      # intermediate out
      snake_case_ =   np.asarray(intermediate_weights[4][0]      )
      snake_case_ =   np.asarray(intermediate_weights[4][1]      )
      set_param(
          torch_block.feed_forward.output.dense      ,      torch.tensor(SCREAMING_SNAKE_CASE      ).transpose(0      ,      1      ).contiguous()      ,      torch.tensor(SCREAMING_SNAKE_CASE      )      ,      )






def  __lowerCAmelCase					(SCREAMING_SNAKE_CASE      ,      SCREAMING_SNAKE_CASE      ,      SCREAMING_SNAKE_CASE      )->						str:




      """simple docstring"""

      snake_case_ =   torch_model.reformer

      # word embeds
      snake_case_ =   np.asarray(weights[1]      )
      set_param(
          torch_model_reformer.embeddings.word_embeddings      ,      torch.tensor(SCREAMING_SNAKE_CASE      )      ,      )

      if isinstance(weights[3]      ,      SCREAMING_SNAKE_CASE      ):
            snake_case_ =   torch_model_reformer.embeddings.position_embeddings
            for emb_idx in range(len(position_embeddings.weights      )      ):
                  snake_case_ =   np.asarray(weights[3][emb_idx][0]      )
                  assert (
                      position_embeddings.weights[emb_idx].shape == emb_weights.shape
                  ), f'''{position_embeddings[emb_idx]} emb does not match'''
                  snake_case_ =   nn.Parameter(torch.tensor(SCREAMING_SNAKE_CASE      )      )

      snake_case_ =   weights[5]
      assert len(torch_model_reformer.encoder.layers      ) * 4 == len(
          SCREAMING_SNAKE_CASE      ), "HF and trax model do not have the same number of layers"
      for layer_idx, layer in enumerate(torch_model_reformer.encoder.layers      ):
            snake_case_ =   trax_layer_weights[4 * layer_idx : 4 * (layer_idx + 1)]
            set_block_weights_in_torch(SCREAMING_SNAKE_CASE      ,      SCREAMING_SNAKE_CASE      ,      SCREAMING_SNAKE_CASE      )

      # output layer norm
      snake_case_ =   np.asarray(weights[7][0]      )
      snake_case_ =   np.asarray(weights[7][1]      )
      set_param(
          torch_model_reformer.encoder.layer_norm      ,      torch.tensor(SCREAMING_SNAKE_CASE      )      ,      torch.tensor(SCREAMING_SNAKE_CASE      )      ,      )

      # output embeddings
      snake_case_ =   np.asarray(weights[9][0]      )
      snake_case_ =   np.asarray(weights[9][1]      )
      set_param(
          torch_model.lm_head.decoder      ,      torch.tensor(SCREAMING_SNAKE_CASE      ).transpose(0      ,      1      ).contiguous()      ,      torch.tensor(SCREAMING_SNAKE_CASE      )      ,      )






def  __lowerCAmelCase					(SCREAMING_SNAKE_CASE      ,      SCREAMING_SNAKE_CASE      ,      SCREAMING_SNAKE_CASE      )->						int:




      """simple docstring"""

      snake_case_ =   ReformerConfig.from_json_file(SCREAMING_SNAKE_CASE      )
      print(f'''Building PyTorch model from configuration: {config}'''      )
      snake_case_ =   ReformerModelWithLMHead(SCREAMING_SNAKE_CASE      )

      with open(SCREAMING_SNAKE_CASE      ,      '''rb'''      ) as f:
            snake_case_ =   pickle.load(SCREAMING_SNAKE_CASE      )['''weights''']

      set_model_weights_in_torch(SCREAMING_SNAKE_CASE      ,      SCREAMING_SNAKE_CASE      ,      config.hidden_size      )

      # Save pytorch-model
      print(f'''Save PyTorch model to {pytorch_dump_path}'''      )
      torch.save(model.state_dict()      ,      SCREAMING_SNAKE_CASE      )


if __name__ == "__main__":
   UpperCAmelCase    =					argparse.ArgumentParser()
   # Required parameters
   parser.add_argument(
       """--trax_model_pkl_path""", default=None, type=str, required=True, help="""Path to the TensorFlow checkpoint path."""
   )
   parser.add_argument(
       """--config_file""",
       default=None,
       type=str,
       required=True,
       help=(
           """The config json file corresponding to the pre-trained Reformer model. \n"""
           """This specifies the model architecture."""
       ),
   )
   parser.add_argument(
       """--pytorch_dump_path""", default=None, type=str, required=True, help="""Path to the output PyTorch model."""
   )
   UpperCAmelCase    =					parser.parse_args()
   convert_trax_checkpoint_to_pytorch(args.trax_model_pkl_path, args.config_file, args.pytorch_dump_path)