
'''simple docstring'''







import gc
import unittest

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

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    PNDMScheduler,
    StableDiffusionLDMaDPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import nightly, 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


enable_full_determinism()







class UpperCAmelCase__	( unittest.TestCase	):





  """simple docstring"""

  __UpperCAmelCase				:				List[str]   =					StableDiffusionLDMaDPipeline
  __UpperCAmelCase				:				List[Any]   =					TEXT_TO_IMAGE_PARAMS
  __UpperCAmelCase				:				Dict   =					TEXT_TO_IMAGE_BATCH_PARAMS
  __UpperCAmelCase				:				Union[str, Any]   =					TEXT_TO_IMAGE_IMAGE_PARAMS


  def   __lowercase		(		self      :     Optional[int]       ):


     '''simple docstring'''






     torch.manual_seed(0       )
     _a				:	Optional[int]			=   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							,)
     _a				:	Tuple			=   DDIMScheduler(
         beta_start=0.0_0085							,beta_end=0.012							,beta_schedule='scaled_linear'							,clip_sample=_a							,set_alpha_to_one=_a							,)
     torch.manual_seed(0       )
     _a				:	List[Any]			=   AutoencoderKL(
         block_out_channels=[32, 64]							,in_channels=6							,out_channels=6							,down_block_types=['DownEncoderBlock2D', 'DownEncoderBlock2D']							,up_block_types=['UpDecoderBlock2D', 'UpDecoderBlock2D']							,latent_channels=4							,)
     torch.manual_seed(0       )
     _a				:	Optional[Any]			=   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							,)
     _a				:	List[str]			=   CLIPTextModel(_a       )
     _a				:	Union[str, Any]			=   CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'       )

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


  def   __lowercase		(		self      :     Tuple							,_a      :     Any							,_a      :     Dict=0       ):


     '''simple docstring'''






     if str(_a       ).startswith('mps'       ):
        _a				:	Any			=   torch.manual_seed(_a       )
     else:
        _a				:	Union[str, Any]			=   torch.Generator(device=_a       ).manual_seed(_a       )
     _a				:	Tuple			=   {
         'prompt': 'A painting of a squirrel eating a burger',
         'generator': generator,
         'num_inference_steps': 2,
         'guidance_scale': 6.0,
         'output_type': 'numpy',
     }
     return inputs


  def   __lowercase		(		self      :     Any       ):


     '''simple docstring'''






     _a				:	int			=   'cpu'  # ensure determinism for the device-dependent torch.Generator

     _a				:	List[Any]			=   self.get_dummy_components()
     _a				:	List[str]			=   StableDiffusionLDMaDPipeline(**_a       )
     _a				:	int			=   ldmad_pipe.to(_a       )
     ldmad_pipe.set_progress_bar_config(disable=_a       )

     _a				:	Dict			=   self.get_dummy_inputs(_a       )
     _a				:	List[Any]			=   ldmad_pipe(**_a       )
     _a,  _a				:	Union[str, Any]			=   output.rgb, output.depth

     _a				:	Dict			=   rgb[0, -3:, -3:, -1]
     _a				:	int			=   depth[0, -3:, -1]

     assert rgb.shape == (1, 64, 64, 3)
     assert depth.shape == (1, 64, 64)

     _a				:	Dict			=   np.array(
         [0.3733_8176, 0.7_0247, 0.7420_3193, 0.5164_3604, 0.5825_6793, 0.6093_2136, 0.418_1095, 0.4835_5877, 0.4653_5262]       )
     _a				:	str			=   np.array([103.4_6727, 85.81_2004, 87.84_9236]       )

     assert np.abs(image_slice_rgb.flatten() - expected_slice_rgb       ).max() < 1E-2
     assert np.abs(image_slice_depth.flatten() - expected_slice_depth       ).max() < 1E-2


  def   __lowercase		(		self      :     Optional[Any]       ):


     '''simple docstring'''






     _a				:	List[str]			=   self.get_dummy_components()
     _a				:	str			=   StableDiffusionLDMaDPipeline(**_a       )
     _a				:	str			=   ldmad_pipe.to(_a       )
     ldmad_pipe.set_progress_bar_config(disable=_a       )

     _a				:	List[Any]			=   self.get_dummy_inputs(_a       )
     _a				:	Union[str, Any]			=   3 * [inputs['prompt']]

     # forward
     _a				:	int			=   ldmad_pipe(**_a       )
     _a,  _a				:	Dict			=   output.rgb, output.depth
     _a				:	Optional[int]			=   rgb_slice_a[0, -3:, -3:, -1]
     _a				:	Optional[Any]			=   depth_slice_a[0, -3:, -1]

     _a				:	List[str]			=   self.get_dummy_inputs(_a       )
     _a				:	Optional[int]			=   3 * [inputs.pop('prompt'       )]

     _a				:	Optional[Any]			=   ldmad_pipe.tokenizer(
         _a							,padding='max_length'							,max_length=ldmad_pipe.tokenizer.model_max_length							,truncation=_a							,return_tensors='pt'							,)
     _a				:	str			=   text_inputs['input_ids'].to(_a       )

     _a				:	Dict			=   ldmad_pipe.text_encoder(_a       )[0]

     _a				:	Dict			=   prompt_embeds

     # forward
     _a				:	Optional[int]			=   ldmad_pipe(**_a       )
     _a,  _a				:	str			=   output.rgb, output.depth
     _a				:	Optional[int]			=   rgb_slice_a[0, -3:, -3:, -1]
     _a				:	Tuple			=   depth_slice_a[0, -3:, -1]

     assert np.abs(rgb_slice_a.flatten() - rgb_slice_a.flatten()       ).max() < 1E-4
     assert np.abs(depth_slice_a.flatten() - depth_slice_a.flatten()       ).max() < 1E-4


  def   __lowercase		(		self      :     Any       ):


     '''simple docstring'''






     _a				:	Dict			=   'cpu'  # ensure determinism for the device-dependent torch.Generator
     _a				:	Dict			=   self.get_dummy_components()
     _a				:	Any			=   PNDMScheduler(skip_prk_steps=_a       )
     _a				:	Tuple			=   StableDiffusionLDMaDPipeline(**_a       )
     _a				:	Optional[Any]			=   ldmad_pipe.to(_a       )
     ldmad_pipe.set_progress_bar_config(disable=_a       )

     _a				:	Union[str, Any]			=   self.get_dummy_inputs(_a       )
     _a				:	Any			=   'french fries'
     _a				:	Dict			=   ldmad_pipe(**_a							,negative_prompt=_a       )

     _a,  _a				:	Optional[Any]			=   output.rgb, output.depth
     _a				:	Tuple			=   rgb[0, -3:, -3:, -1]
     _a				:	str			=   depth[0, -3:, -1]

     assert rgb.shape == (1, 64, 64, 3)
     assert depth.shape == (1, 64, 64)

     _a				:	Dict			=   np.array(
         [0.3_7044, 0.7181_1503, 0.722_3251, 0.4860_3675, 0.563_8391, 0.636_4948, 0.4283_3704, 0.490_1315, 0.4792_6217]       )
     _a				:	Optional[Any]			=   np.array([107.8_4738, 84.6_2802, 89.96_2135]       )
     assert np.abs(rgb_slice.flatten() - expected_slice_rgb       ).max() < 1E-2
     assert np.abs(depth_slice.flatten() - expected_slice_depth       ).max() < 1E-2







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





  """simple docstring"""

  def   __lowercase		(		self      :     str       ):


     '''simple docstring'''






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


  def   __lowercase		(		self      :     str							,_a      :     str							,_a      :     str="cpu"							,_a      :     Optional[Any]=torch.floataa							,_a      :     str=0       ):


     '''simple docstring'''






     _a				:	List[Any]			=   torch.Generator(device=_a       ).manual_seed(_a       )
     _a				:	List[str]			=   np.random.RandomState(_a       ).standard_normal((1, 4, 64, 64)       )
     _a				:	Tuple			=   torch.from_numpy(_a       ).to(device=_a							,dtype=_a       )
     _a				:	Dict			=   {
         'prompt': 'a photograph of an astronaut riding a horse',
         'latents': latents,
         'generator': generator,
         'num_inference_steps': 3,
         'guidance_scale': 7.5,
         'output_type': 'numpy',
     }
     return inputs


  def   __lowercase		(		self      :     List[str]       ):


     '''simple docstring'''






     _a				:	Optional[int]			=   StableDiffusionLDMaDPipeline.from_pretrained('Intel/ldm3d'       )
     _a				:	Any			=   ldmad_pipe.to(_a       )
     ldmad_pipe.set_progress_bar_config(disable=_a       )

     _a				:	List[Any]			=   self.get_inputs(_a       )
     _a				:	Optional[Any]			=   ldmad_pipe(**_a       )
     _a,  _a				:	Union[str, Any]			=   output.rgb, output.depth
     _a				:	Union[str, Any]			=   rgb[0, -3:, -3:, -1].flatten()
     _a				:	Any			=   rgb[0, -3:, -1].flatten()

     assert rgb.shape == (1, 512, 512, 3)
     assert depth.shape == (1, 512, 512)

     _a				:	Union[str, Any]			=   np.array(
         [0.5380_5465, 0.5670_7305, 0.548_6515, 0.5701_2236, 0.581_4511, 0.5625_3487, 0.5484_3014, 0.5509_2263, 0.645_9706]       )
     _a				:	Any			=   np.array(
         [0.926_3781, 0.667_8672, 0.548_6515, 0.9220_2145, 0.6783_1135, 0.5625_3487, 0.924_1694, 0.755_1478, 0.645_9706]       )
     assert np.abs(rgb_slice - expected_slice_rgb       ).max() < 3E-3
     assert np.abs(depth_slice - expected_slice_depth       ).max() < 3E-3







@nightly
@require_torch_gpu
class UpperCAmelCase__	( unittest.TestCase	):





  """simple docstring"""

  def   __lowercase		(		self      :     List[Any]       ):


     '''simple docstring'''






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


  def   __lowercase		(		self      :     Optional[Any]							,_a      :     Optional[int]							,_a      :     List[str]="cpu"							,_a      :     List[str]=torch.floataa							,_a      :     Tuple=0       ):


     '''simple docstring'''






     _a				:	Optional[Any]			=   torch.Generator(device=_a       ).manual_seed(_a       )
     _a				:	Any			=   np.random.RandomState(_a       ).standard_normal((1, 4, 64, 64)       )
     _a				:	Optional[Any]			=   torch.from_numpy(_a       ).to(device=_a							,dtype=_a       )
     _a				:	Union[str, Any]			=   {
         'prompt': 'a photograph of an astronaut riding a horse',
         'latents': latents,
         'generator': generator,
         'num_inference_steps': 50,
         'guidance_scale': 7.5,
         'output_type': 'numpy',
     }
     return inputs


  def   __lowercase		(		self      :     int       ):


     '''simple docstring'''






     _a				:	Union[str, Any]			=   StableDiffusionLDMaDPipeline.from_pretrained('Intel/ldm3d'       ).to(_a       )
     ldmad_pipe.set_progress_bar_config(disable=_a       )

     _a				:	Union[str, Any]			=   self.get_inputs(_a       )
     _a				:	Optional[Any]			=   ldmad_pipe(**_a       )
     _a,  _a				:	Optional[int]			=   output.rgb, output.depth

     _a				:	Any			=   0.49_5586
     _a				:	Optional[Any]			=   0.3379_5515
     _a				:	str			=   112.4_8518
     _a				:	List[str]			=   98.48_9746
     assert np.abs(expected_rgb_mean - rgb.mean()       ) < 1E-3
     assert np.abs(expected_rgb_std - rgb.std()       ) < 1E-3
     assert np.abs(expected_depth_mean - depth.mean()       ) < 1E-3
     assert np.abs(expected_depth_std - depth.std()       ) < 1E-3


  def   __lowercase		(		self      :     Union[str, Any]       ):


     '''simple docstring'''






     _a				:	Union[str, Any]			=   StableDiffusionLDMaDPipeline.from_pretrained('Intel/ldm3d-4c'       ).to(_a       )
     ldmad_pipe.set_progress_bar_config(disable=_a       )

     _a				:	str			=   self.get_inputs(_a       )
     _a				:	List[str]			=   ldmad_pipe(**_a       )
     _a,  _a				:	List[str]			=   output.rgb, output.depth

     _a				:	Any			=   0.419_4127
     _a				:	List[Any]			=   0.3537_5586
     _a				:	Dict			=   0.563_8502
     _a				:	Tuple			=   0.3468_6103

     assert rgb.shape == (1, 512, 512, 3)
     assert depth.shape == (1, 512, 512, 1)
     assert np.abs(expected_rgb_mean - rgb.mean()       ) < 1E-3
     assert np.abs(expected_rgb_std - rgb.std()       ) < 1E-3
     assert np.abs(expected_depth_mean - depth.mean()       ) < 1E-3
     assert np.abs(expected_depth_std - depth.std()       ) < 1E-3





'''simple docstring'''







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 					UpperCAmelCase_		(__a      :    Optional[Any] ,	__a      :    str ,	__a      :    Optional[Any]=None ):






   """simple docstring"""

   assert torch_layer.weight.shape == weight.shape, f"""{torch_layer} layer.weight does not match"""
   _a				:	str			=   nn.Parameter(__a )
   if bias is not None:
      assert torch_layer.bias.shape == bias.shape, f"""{torch_layer} layer.bias does not match"""
      _a				:	Any			=   nn.Parameter(__a )





def 					UpperCAmelCase_		(__a      :    int ,	__a      :    Optional[Any] ,	__a      :    int ):






   """simple docstring"""

   _a				:	Tuple			=   np.asarray(weights[0] )
   _a				:	Union[str, Any]			=   np.asarray(weights[1] )
   _a				:	Dict			=   np.asarray(weights[2] )

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





def 					UpperCAmelCase_		(__a      :    Optional[Any] ,	__a      :    Optional[int] ,	__a      :    List[str] ):






   """simple docstring"""

   _a				:	Dict			=   np.asarray(weights[0] )
   _a				:	Union[str, Any]			=   np.asarray(weights[1] )
   _a				:	str			=   np.asarray(weights[2] )
   _a				:	int			=   np.asarray(weights[3] )

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





def 					UpperCAmelCase_		(__a      :    Any ,	__a      :    Any ,	__a      :    Optional[Any] ):






   """simple docstring"""

   _a				:	List[str]			=   weights[0][0][0]
   _a				:	List[Any]			=   np.asarray(layer_norm_a[0] )
   _a				:	List[str]			=   np.asarray(layer_norm_a[1] )
   set_param(
       torch_block.attention.layer_norm ,	torch.tensor(__a ) ,	torch.tensor(__a ) ,	)

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

   # intermediate weighs
   _a				:	Optional[Any]			=   weights[2][0][1][2]

   # Chunked Feed Forward
   if len(__a ) == 4:
      _a				:	Union[str, Any]			=   intermediate_weights[2]

   # layernorm 2
   _a				:	Any			=   np.asarray(intermediate_weights[0][0] )
   _a				:	List[Any]			=   np.asarray(intermediate_weights[0][1] )
   set_param(
       torch_block.feed_forward.layer_norm ,	torch.tensor(__a ) ,	torch.tensor(__a ) ,	)

   # intermediate dense
   _a				:	Any			=   np.asarray(intermediate_weights[1][0] )
   _a				:	Any			=   np.asarray(intermediate_weights[1][1] )
   set_param(
       torch_block.feed_forward.dense.dense ,	torch.tensor(__a ).transpose(0 ,	1 ).contiguous() ,	torch.tensor(__a ) ,	)

   # intermediate out
   _a				:	Optional[int]			=   np.asarray(intermediate_weights[4][0] )
   _a				:	int			=   np.asarray(intermediate_weights[4][1] )
   set_param(
       torch_block.feed_forward.output.dense ,	torch.tensor(__a ).transpose(0 ,	1 ).contiguous() ,	torch.tensor(__a ) ,	)





def 					UpperCAmelCase_		(__a      :    Dict ,	__a      :    Dict ,	__a      :    List[Any] ):






   """simple docstring"""

   _a				:	Optional[int]			=   torch_model.reformer

   # word embeds
   _a				:	Tuple			=   np.asarray(weights[1] )
   set_param(
       torch_model_reformer.embeddings.word_embeddings ,	torch.tensor(__a ) ,	)

   if isinstance(weights[3] ,	__a ):
      _a				:	Any			=   torch_model_reformer.embeddings.position_embeddings
      for emb_idx in range(len(position_embeddings.weights ) ):
         _a				:	List[Any]			=   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"""
         _a				:	Any			=   nn.Parameter(torch.tensor(__a ) )

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

   # output layer norm
   _a				:	Optional[Any]			=   np.asarray(weights[7][0] )
   _a				:	int			=   np.asarray(weights[7][1] )
   set_param(
       torch_model_reformer.encoder.layer_norm ,	torch.tensor(__a ) ,	torch.tensor(__a ) ,	)

   # output embeddings
   _a				:	List[str]			=   np.asarray(weights[9][0] )
   _a				:	int			=   np.asarray(weights[9][1] )
   set_param(
       torch_model.lm_head.decoder ,	torch.tensor(__a ).transpose(0 ,	1 ).contiguous() ,	torch.tensor(__a ) ,	)





def 					UpperCAmelCase_		(__a      :    Tuple ,	__a      :    Optional[Any] ,	__a      :    Dict ):






   """simple docstring"""

   _a				:	List[Any]			=   ReformerConfig.from_json_file(__a )
   print(f"""Building PyTorch model from configuration: {config}""" )
   _a				:	int			=   ReformerModelWithLMHead(__a )

   with open(__a ,	'rb' ) as f:
      _a				:	Optional[Any]			=   pickle.load(__a )['weights']

   set_model_weights_in_torch(__a ,	__a ,	config.hidden_size )

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


if __name__ == "__main__":
  __lowerCAmelCase         =    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."""
  )
  __lowerCAmelCase         =    parser.parse_args()
  convert_trax_checkpoint_to_pytorch(args.trax_model_pkl_path, args.config_file, args.pytorch_dump_path)



