




'''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 						snake_case_       (UpperCamelCase			:    Optional[Any]					,		UpperCamelCase			:    int					,		UpperCamelCase			:    Union[str, Any]=None ):





							'''simple docstring'''


							assert torch_layer.weight.shape == weight.shape, f'{torch_layer} layer.weight does not match'
							_a      	=							nn.Parameter(UpperCamelCase )
							if bias is not None:
														assert torch_layer.bias.shape == bias.shape, f'{torch_layer} layer.bias does not match'
														_a      	=							nn.Parameter(UpperCamelCase )
def 						snake_case_       (UpperCamelCase			:    Dict					,		UpperCamelCase			:    Optional[int]					,		UpperCamelCase			:    Optional[int] ):





							'''simple docstring'''


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

							set_param(
							    torch_layer.self_attention.query_key					,		torch.tensor(UpperCamelCase ).transpose(1					,		2 ).contiguous().view(-1					,		UpperCamelCase )					,		)
							set_param(
							    torch_layer.self_attention.value					,		torch.tensor(UpperCamelCase ).transpose(1					,		2 ).contiguous().view(-1					,		UpperCamelCase )					,		)
							set_param(
							    torch_layer.output.dense					,		torch.tensor(UpperCamelCase ).view(-1					,		UpperCamelCase ).contiguous().transpose(0					,		1 )					,		)
def 						snake_case_       (UpperCamelCase			:    List[Any]					,		UpperCamelCase			:    Union[str, Any]					,		UpperCamelCase			:    int ):





							'''simple docstring'''


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

							set_param(
							    torch_layer.self_attention.query					,		torch.tensor(UpperCamelCase ).transpose(1					,		2 ).contiguous().view(-1					,		UpperCamelCase )					,		)
							set_param(
							    torch_layer.self_attention.key					,		torch.tensor(UpperCamelCase ).transpose(1					,		2 ).contiguous().view(-1					,		UpperCamelCase )					,		)
							set_param(
							    torch_layer.self_attention.value					,		torch.tensor(UpperCamelCase ).transpose(1					,		2 ).contiguous().view(-1					,		UpperCamelCase )					,		)
							set_param(
							    torch_layer.output.dense					,		torch.tensor(UpperCamelCase ).view(-1					,		UpperCamelCase ).contiguous().transpose(0					,		1 )					,		)
def 						snake_case_       (UpperCamelCase			:    int					,		UpperCamelCase			:    Any					,		UpperCamelCase			:    int ):





							'''simple docstring'''


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

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

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

							# Chunked Feed Forward
							if len(UpperCamelCase ) == 4:
														_a      	=							intermediate_weights[2]

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

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

							# intermediate out
							_a      	=							np.asarray(intermediate_weights[4][0] )
							_a      	=							np.asarray(intermediate_weights[4][1] )
							set_param(
							    torch_block.feed_forward.output.dense					,		torch.tensor(UpperCamelCase ).transpose(0					,		1 ).contiguous()					,		torch.tensor(UpperCamelCase )					,		)
def 						snake_case_       (UpperCamelCase			:    Dict					,		UpperCamelCase			:    str					,		UpperCamelCase			:    List[Any] ):





							'''simple docstring'''


							_a      	=							torch_model.reformer

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

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

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

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

							# output embeddings
							_a      	=							np.asarray(weights[9][0] )
							_a      	=							np.asarray(weights[9][1] )
							set_param(
							    torch_model.lm_head.decoder					,		torch.tensor(UpperCamelCase ).transpose(0					,		1 ).contiguous()					,		torch.tensor(UpperCamelCase )					,		)
def 						snake_case_       (UpperCamelCase			:    str					,		UpperCamelCase			:    Optional[Any]					,		UpperCamelCase			:    Tuple ):





							'''simple docstring'''


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

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

							set_model_weights_in_torch(UpperCamelCase					,		UpperCamelCase					,		config.hidden_size )

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


if __name__ == "__main__":
				_snake_case : List[str]  				=						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.'
				)
				_snake_case : Tuple  				=						parser.parse_args()
				convert_trax_checkpoint_to_pytorch(args.trax_model_pkl_path, args.config_file, args.pytorch_dump_path)




import os
import tempfile
import unittest

import numpy as np

from diffusers.utils import is_flax_available
from diffusers.utils.testing_utils import require_flax, slow


if is_flax_available():
	import jax
	import jax.numpy as jnp
	from flax.jax_utils import replicate
	from flax.training.common_utils import shard

	from diffusers import FlaxDDIMScheduler, FlaxDiffusionPipeline, FlaxStableDiffusionPipeline



@require_flax
class        _a		(unittest.TestCase							):


					'''simple docstring'''






					def     __A     (					self     ):
						with tempfile.TemporaryDirectory() as tmpdirname:
							# pipeline has Flax weights
							A__				:       Optional[int]	  =						FlaxDiffusionPipeline.from_pretrained(
							    """hf-internal-testing/tiny-stable-diffusion-pipe"""       ,			safety_checker=A__       ,			cache_dir=A__     )

							A__				:       int	  =						[t[-1] for t in os.walk(os.path.join(A__       ,			os.listdir(A__     )[0]       ,			"""snapshots"""     )     )]
							A__				:       str	  =						[item for sublist in all_root_files for item in sublist]

							# None of the downloaded files should be a PyTorch file even if we have some here:
							# https://huggingface.co/hf-internal-testing/tiny-stable-diffusion-pipe/blob/main/unet/diffusion_pytorch_model.bin
							assert not any(f.endswith(""".bin"""     ) for f in files     )



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


					'''simple docstring'''






					def     __A     (					self     ):
						A__  ,  A__				:       Any	  =						FlaxStableDiffusionPipeline.from_pretrained(
						    """hf-internal-testing/tiny-stable-diffusion-pipe"""       ,			safety_checker=A__     )

						A__				:       Optional[int]	  =						(
						    """A cinematic film still of Morgan Freeman starring as Jimi Hendrix, portrait, 40mm lens, shallow depth of"""
						    """ field, close up, split lighting, cinematic"""
						)

						A__				:       Tuple	  =						jax.random.PRNGKey(0     )
						A__				:       int	  =						4

						A__				:       Optional[Any]	  =						jax.device_count()
						A__				:       Union[str, Any]	  =						num_samples * [prompt]
						A__				:       str	  =						pipeline.prepare_inputs(A__     )

						# shard inputs and rng
						A__				:       str	  =						replicate(A__     )
						A__				:       Any	  =						jax.random.split(A__       ,			A__     )
						A__				:       Optional[Any]	  =						shard(A__     )

						A__				:       int	  =						pipeline(A__       ,			A__       ,			A__       ,			A__       ,			jit=A__     ).images

						assert images.shape == (num_samples, 1, 64, 64, 3)
						if jax.device_count() == 8:
							assert np.abs(np.abs(images[0, 0, :2, :2, -2:]       ,			dtype=np.floataa     ).sum() - 4.1_5_1_4_7_4_5     ) < 1e-3
							assert np.abs(np.abs(A__       ,			dtype=np.floataa     ).sum() - 4_9_9_4_7.8_7_5     ) < 5e-1

						A__				:       str	  =						pipeline.numpy_to_pil(np.asarray(images.reshape((num_samples,) + images.shape[-3:]     )     )     )
						assert len(A__     ) == num_samples






					def     __A     (					self     ):
						A__  ,  A__				:       Any	  =						FlaxStableDiffusionPipeline.from_pretrained(
						    """CompVis/stable-diffusion-v1-4"""       ,			revision="""flax"""       ,			safety_checker=A__     )

						A__				:       Tuple	  =						(
						    """A cinematic film still of Morgan Freeman starring as Jimi Hendrix, portrait, 40mm lens, shallow depth of"""
						    """ field, close up, split lighting, cinematic"""
						)

						A__				:       Tuple	  =						jax.random.PRNGKey(0     )
						A__				:       Optional[Any]	  =						50

						A__				:       List[str]	  =						jax.device_count()
						A__				:       Dict	  =						num_samples * [prompt]
						A__				:       Union[str, Any]	  =						pipeline.prepare_inputs(A__     )

						# shard inputs and rng
						A__				:       int	  =						replicate(A__     )
						A__				:       List[str]	  =						jax.random.split(A__       ,			A__     )
						A__				:       Optional[Any]	  =						shard(A__     )

						A__				:       Optional[Any]	  =						pipeline(A__       ,			A__       ,			A__       ,			A__       ,			jit=A__     ).images

						assert images.shape == (num_samples, 1, 512, 512, 3)
						if jax.device_count() == 8:
							assert np.abs((np.abs(images[0, 0, :2, :2, -2:]       ,			dtype=np.floataa     ).sum() - 0.0_5_6_5_2_4_0_1)     ) < 1e-3
							assert np.abs((np.abs(A__       ,			dtype=np.floataa     ).sum() - 2_3_8_3_8_0_8.2)     ) < 5e-1






					def     __A     (					self     ):
						A__  ,  A__				:       int	  =						FlaxStableDiffusionPipeline.from_pretrained(
						    """CompVis/stable-diffusion-v1-4"""       ,			revision="""bf16"""       ,			dtype=jnp.bfloataa       ,			safety_checker=A__     )

						A__				:       Any	  =						(
						    """A cinematic film still of Morgan Freeman starring as Jimi Hendrix, portrait, 40mm lens, shallow depth of"""
						    """ field, close up, split lighting, cinematic"""
						)

						A__				:       Dict	  =						jax.random.PRNGKey(0     )
						A__				:       str	  =						50

						A__				:       Any	  =						jax.device_count()
						A__				:       List[Any]	  =						num_samples * [prompt]
						A__				:       Any	  =						pipeline.prepare_inputs(A__     )

						# shard inputs and rng
						A__				:       Any	  =						replicate(A__     )
						A__				:       str	  =						jax.random.split(A__       ,			A__     )
						A__				:       Optional[Any]	  =						shard(A__     )

						A__				:       Dict	  =						pipeline(A__       ,			A__       ,			A__       ,			A__       ,			jit=A__     ).images

						assert images.shape == (num_samples, 1, 512, 512, 3)
						if jax.device_count() == 8:
							assert np.abs((np.abs(images[0, 0, :2, :2, -2:]       ,			dtype=np.floataa     ).sum() - 0.0_4_0_0_3_9_0_6)     ) < 1e-3
							assert np.abs((np.abs(A__       ,			dtype=np.floataa     ).sum() - 2_3_7_3_5_1_6.7_5)     ) < 5e-1






					def     __A     (					self     ):
						A__  ,  A__				:       Tuple	  =						FlaxStableDiffusionPipeline.from_pretrained(
						    """CompVis/stable-diffusion-v1-4"""       ,			revision="""bf16"""       ,			dtype=jnp.bfloataa     )

						A__				:       List[str]	  =						(
						    """A cinematic film still of Morgan Freeman starring as Jimi Hendrix, portrait, 40mm lens, shallow depth of"""
						    """ field, close up, split lighting, cinematic"""
						)

						A__				:       int	  =						jax.random.PRNGKey(0     )
						A__				:       Optional[int]	  =						50

						A__				:       Optional[Any]	  =						jax.device_count()
						A__				:       Optional[int]	  =						num_samples * [prompt]
						A__				:       Optional[int]	  =						pipeline.prepare_inputs(A__     )

						# shard inputs and rng
						A__				:       Optional[int]	  =						replicate(A__     )
						A__				:       Tuple	  =						jax.random.split(A__       ,			A__     )
						A__				:       Optional[int]	  =						shard(A__     )

						A__				:       Optional[int]	  =						pipeline(A__       ,			A__       ,			A__       ,			A__       ,			jit=A__     ).images

						assert images.shape == (num_samples, 1, 512, 512, 3)
						if jax.device_count() == 8:
							assert np.abs((np.abs(images[0, 0, :2, :2, -2:]       ,			dtype=np.floataa     ).sum() - 0.0_4_0_0_3_9_0_6)     ) < 1e-3
							assert np.abs((np.abs(A__       ,			dtype=np.floataa     ).sum() - 2_3_7_3_5_1_6.7_5)     ) < 5e-1






					def     __A     (					self     ):
						A__				:       int	  =						FlaxDDIMScheduler(
						    beta_start=0.0_0_0_8_5       ,			beta_end=0.0_1_2       ,			beta_schedule="""scaled_linear"""       ,			set_alpha_to_one=A__       ,			steps_offset=1       ,			)

						A__  ,  A__				:       Tuple	  =						FlaxStableDiffusionPipeline.from_pretrained(
						    """CompVis/stable-diffusion-v1-4"""       ,			revision="""bf16"""       ,			dtype=jnp.bfloataa       ,			scheduler=A__       ,			safety_checker=A__       ,			)
						A__				:       Union[str, Any]	  =						scheduler.create_state()

						A__				:       Optional[Any]	  =						scheduler_state

						A__				:       Union[str, Any]	  =						(
						    """A cinematic film still of Morgan Freeman starring as Jimi Hendrix, portrait, 40mm lens, shallow depth of"""
						    """ field, close up, split lighting, cinematic"""
						)

						A__				:       Tuple	  =						jax.random.PRNGKey(0     )
						A__				:       List[Any]	  =						50

						A__				:       Optional[int]	  =						jax.device_count()
						A__				:       Dict	  =						num_samples * [prompt]
						A__				:       List[Any]	  =						pipeline.prepare_inputs(A__     )

						# shard inputs and rng
						A__				:       Tuple	  =						replicate(A__     )
						A__				:       Dict	  =						jax.random.split(A__       ,			A__     )
						A__				:       Dict	  =						shard(A__     )

						A__				:       int	  =						pipeline(A__       ,			A__       ,			A__       ,			A__       ,			jit=A__     ).images

						assert images.shape == (num_samples, 1, 512, 512, 3)
						if jax.device_count() == 8:
							assert np.abs((np.abs(images[0, 0, :2, :2, -2:]       ,			dtype=np.floataa     ).sum() - 0.0_4_5_0_4_3_9_4_5)     ) < 1e-3
							assert np.abs((np.abs(A__       ,			dtype=np.floataa     ).sum() - 2_3_4_7_6_9_3.5)     ) < 5e-1






					def     __A     (					self     ):
						A__				:       Optional[int]	  =						(
						    """A cinematic film still of Morgan Freeman starring as Jimi Hendrix, portrait, 40mm lens, shallow depth of"""
						    """ field, close up, split lighting, cinematic"""
						)

						A__				:       Optional[int]	  =						jax.device_count()
						A__				:       str	  =						num_samples * [prompt]
						A__				:       Any	  =						jax.random.split(jax.random.PRNGKey(0     )       ,			A__     )

						A__  ,  A__				:       Optional[Any]	  =						FlaxStableDiffusionPipeline.from_pretrained(
						    """CompVis/stable-diffusion-v1-4"""       ,			revision="""bf16"""       ,			dtype=jnp.bfloataa       ,			safety_checker=A__       ,			)

						A__				:       Tuple	  =						replicate(A__     )
						A__				:       Dict	  =						pipeline.prepare_inputs(A__     )
						A__				:       str	  =						shard(A__     )
						A__				:       Tuple	  =						pipeline(A__       ,			A__       ,			A__       ,			jit=A__     ).images
						assert images.shape == (num_samples, 1, 512, 512, 3)
						A__				:       Optional[Any]	  =						images[2, 0, 256, 10:17, 1]

						# With memory efficient attention
						A__  ,  A__				:       Any	  =						FlaxStableDiffusionPipeline.from_pretrained(
						    """CompVis/stable-diffusion-v1-4"""       ,			revision="""bf16"""       ,			dtype=jnp.bfloataa       ,			safety_checker=A__       ,			use_memory_efficient_attention=A__       ,			)

						A__				:       Any	  =						replicate(A__     )
						A__				:       Optional[int]	  =						pipeline.prepare_inputs(A__     )
						A__				:       Dict	  =						shard(A__     )
						A__				:       Tuple	  =						pipeline(A__       ,			A__       ,			A__       ,			jit=A__     ).images
						assert images_eff.shape == (num_samples, 1, 512, 512, 3)
						A__				:       Union[str, Any]	  =						images[2, 0, 256, 10:17, 1]

						# I checked the results visually and they are very similar. However, I saw that the max diff is `1` and the `sum`
						# over the 8 images is exactly `256`, which is very suspicious. Testing a random slice for now.
						assert abs(slice_eff - slice     ).max() < 1e-2



