
import random
import unittest

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

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    UNetaDConditionModel,
    VideoToVideoSDPipeline,
)
from diffusers.utils import floats_tensor, is_xformers_available, skip_mps
from diffusers.utils.testing_utils import enable_full_determinism, slow, torch_device

from ..pipeline_params import (
    TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS,
    TEXT_GUIDED_IMAGE_VARIATION_PARAMS,
)
from ..test_pipelines_common import PipelineTesterMixin


enable_full_determinism()




@skip_mps
class    SCREAMING_SNAKE_CASE__     (	UpperCAmelCase							,unittest.TestCase					):





	'''simple docstring'''



	_UpperCAmelCase       :		Optional[Any]       =					VideoToVideoSDPipeline
	_UpperCAmelCase       :		Union[str, Any]       =					TEXT_GUIDED_IMAGE_VARIATION_PARAMS.union({"video"}					) - {"image", "width", "height"}
	_UpperCAmelCase       :		Optional[Any]       =					TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS.union({"video"}					) - {"image"}
	_UpperCAmelCase       :		Union[str, Any]       =					PipelineTesterMixin.required_optional_params - {"latents"}
	_UpperCAmelCase       :		str       =					False

	# No `output_type`.
	_UpperCAmelCase       :		List[str]       =					frozenset(
	    [
	        "num_inference_steps",
	        "generator",
	        "latents",
	        "return_dict",
	        "callback",
	        "callback_steps",
	    ]					)



	def  A							(  self	:    Optional[int]	):


				'''simple docstring'''


				torch.manual_seed(0	)
				_snake_case		=			UNetaDConditionModel(
				    block_out_channels=(32, 64, 64, 64)			,  layers_per_block=2			,  sample_size=32			,  in_channels=4			,  out_channels=4			,  down_block_types=('CrossAttnDownBlock3D', 'CrossAttnDownBlock3D', 'CrossAttnDownBlock3D', 'DownBlock3D')			,  up_block_types=('UpBlock3D', 'CrossAttnUpBlock3D', 'CrossAttnUpBlock3D', 'CrossAttnUpBlock3D')			,  cross_attention_dim=32			,  attention_head_dim=4			,  )
				_snake_case		=			DDIMScheduler(
				    beta_start=0.00085			,  beta_end=0.012			,  beta_schedule='scaled_linear'			,  clip_sample=lowercase			,  set_alpha_to_one=lowercase			,  )
				torch.manual_seed(0	)
				_snake_case		=			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			,  sample_size=128			,  )
				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=1_000			,  hidden_act='gelu'			,  projection_dim=512			,  )
				_snake_case		=			CLIPTextModel(lowercase	)
				_snake_case		=			CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'	)

				_snake_case		=			{
				    'unet': unet,
				    'scheduler': scheduler,
				    'vae': vae,
				    'text_encoder': text_encoder,
				    'tokenizer': tokenizer,
				}
				return components



	def  A							(  self	:    str			,  lowercase	:    str			,  lowercase	:    List[str]=0	):


				'''simple docstring'''


				_snake_case		=			floats_tensor((1, 3, 3, 32, 32)			,  rng=random.Random(lowercase	)	).to(lowercase	)

				if str(lowercase	).startswith('mps'	):
							_snake_case		=			torch.manual_seed(lowercase	)
				else:
							_snake_case		=			torch.Generator(device=lowercase	).manual_seed(lowercase	)
				_snake_case		=			{
				    'prompt': 'A painting of a squirrel eating a burger',
				    'video': video,
				    'generator': generator,
				    'num_inference_steps': 2,
				    'guidance_scale': 6.0,
				    'output_type': 'pt',
				}
				return inputs



	def  A							(  self	:    Any	):


				'''simple docstring'''


				_snake_case		=			'cpu'  # ensure determinism for the device-dependent torch.Generator
				_snake_case		=			self.get_dummy_components()
				_snake_case		=			VideoToVideoSDPipeline(**lowercase	)
				_snake_case		=			sd_pipe.to(lowercase	)
				sd_pipe.set_progress_bar_config(disable=lowercase	)

				_snake_case		=			self.get_dummy_inputs(lowercase	)
				_snake_case		=			'np'
				_snake_case		=			sd_pipe(**lowercase	).frames
				_snake_case		=			frames[0][-3:, -3:, -1]

				assert frames[0].shape == (32, 32, 3)
				_snake_case		=			np.array([106, 117, 113, 174, 137, 112, 148, 151, 131]	)

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



	@unittest.skipIf(
	    torch_device != 'cuda' or not is_xformers_available()			,  reason='XFormers attention is only available with CUDA and `xformers` installed'			,  )
	def  A							(  self	:    Tuple	):


				'''simple docstring'''


				self._test_xformers_attention_forwardGenerator_pass(test_mean_pixel_difference=lowercase			,  expected_max_diff=5E-3	)



	@unittest.skip(reason='Batching needs to be properly figured out first for this pipeline.'	)
	def  A							(  self	:    Optional[Any]	):


				'''simple docstring'''


				pass



	@unittest.skip(reason='Batching needs to be properly figured out first for this pipeline.'	)
	def  A							(  self	:    Union[str, Any]	):


				'''simple docstring'''


				pass



	@unittest.skip(reason='`num_images_per_prompt` argument is not supported for this pipeline.'	)
	def  A							(  self	:    Tuple	):


				'''simple docstring'''


				pass





	def  A							(  self	:    Dict	):


				'''simple docstring'''


				return super().test_progress_bar()




@slow
@skip_mps
class    SCREAMING_SNAKE_CASE__     (	unittest.TestCase					):





	'''simple docstring'''



	def  A							(  self	:    List[Any]	):


				'''simple docstring'''


				_snake_case		=			VideoToVideoSDPipeline.from_pretrained('cerspense/zeroscope_v2_XL'			,  torch_dtype=torch.floataa	)
				pipe.enable_model_cpu_offload()

				# 10 frames
				_snake_case		=			torch.Generator(device='cpu'	).manual_seed(0	)
				_snake_case		=			torch.randn((1, 10, 3, 1_024, 576)			,  generator=lowercase	)
				_snake_case		=			video.to('cuda'	)

				_snake_case		=			'Spiderman is surfing'

				_snake_case		=			pipe(lowercase			,  video=lowercase			,  generator=lowercase			,  num_inference_steps=3			,  output_type='pt'	).frames

				_snake_case		=			np.array([-1.0458984, -1.1279297, -0.9663086, -0.91503906, -0.75097656]	)
				assert np.abs(video_frames.cpu().numpy()[0, 0, 0, 0, -5:] - expected_array	).sum() < 1E-2

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        a_		(       __lowercase   :	Dict       ,					__lowercase   :	int       ,					__lowercase   :	Optional[Any]=None ) ->  Any:
			# set parameter of one layer
			assert torch_layer.weight.shape == weight.shape, f'''{torch_layer} layer.weight does not match'''
			_snake_case		=			nn.Parameter(__lowercase )
			if bias is not None:
						assert torch_layer.bias.shape == bias.shape, f'''{torch_layer} layer.bias does not match'''
						_snake_case		=			nn.Parameter(__lowercase )






def        a_		(       __lowercase   :	Any       ,					__lowercase   :	Dict       ,					__lowercase   :	Union[str, Any] ) ->  Optional[Any]:
			# set torch weights for 1-to-1 comparison
			_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(__lowercase ).transpose(1       ,					2 ).contiguous().view(-1       ,					__lowercase )       ,					)
			set_param(
			    torch_layer.self_attention.value       ,					torch.tensor(__lowercase ).transpose(1       ,					2 ).contiguous().view(-1       ,					__lowercase )       ,					)
			set_param(
			    torch_layer.output.dense       ,					torch.tensor(__lowercase ).view(-1       ,					__lowercase ).contiguous().transpose(0       ,					1 )       ,					)






def        a_		(       __lowercase   :	str       ,					__lowercase   :	Tuple       ,					__lowercase   :	Any ) ->  Optional[Any]:
			# set torch weights for 1-to-1 comparison
			_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(__lowercase ).transpose(1       ,					2 ).contiguous().view(-1       ,					__lowercase )       ,					)
			set_param(
			    torch_layer.self_attention.key       ,					torch.tensor(__lowercase ).transpose(1       ,					2 ).contiguous().view(-1       ,					__lowercase )       ,					)
			set_param(
			    torch_layer.self_attention.value       ,					torch.tensor(__lowercase ).transpose(1       ,					2 ).contiguous().view(-1       ,					__lowercase )       ,					)
			set_param(
			    torch_layer.output.dense       ,					torch.tensor(__lowercase ).view(-1       ,					__lowercase ).contiguous().transpose(0       ,					1 )       ,					)






def        a_		(       __lowercase   :	Dict       ,					__lowercase   :	List[str]       ,					__lowercase   :	Union[str, Any] ) ->  Optional[Any]:
			# layernorm 1
			_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(__lowercase )       ,					torch.tensor(__lowercase )       ,					)

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

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

			# Chunked Feed Forward
			if len(__lowercase ) == 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(__lowercase )       ,					torch.tensor(__lowercase )       ,					)

			# 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(__lowercase ).transpose(0       ,					1 ).contiguous()       ,					torch.tensor(__lowercase )       ,					)

			# 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(__lowercase ).transpose(0       ,					1 ).contiguous()       ,					torch.tensor(__lowercase )       ,					)






def        a_		(       __lowercase   :	Tuple       ,					__lowercase   :	Tuple       ,					__lowercase   :	Dict ) ->  Optional[int]:
			# reformer model
			_snake_case		=			torch_model.reformer

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

			if isinstance(weights[3]       ,					__lowercase ):
						_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(__lowercase ) )

			_snake_case		=			weights[5]
			assert len(torch_model_reformer.encoder.layers ) * 4 == len(
			    __lowercase ), "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(__lowercase       ,					__lowercase       ,					__lowercase )

			# 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(__lowercase )       ,					torch.tensor(__lowercase )       ,					)

			# 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(__lowercase ).transpose(0       ,					1 ).contiguous()       ,					torch.tensor(__lowercase )       ,					)






def        a_		(       __lowercase   :	Optional[Any]       ,					__lowercase   :	Dict       ,					__lowercase   :	List[Any] ) ->  Optional[int]:
			# Initialise PyTorch model
			_snake_case		=			ReformerConfig.from_json_file(__lowercase )
			print(f'''Building PyTorch model from configuration: {config}''' )
			_snake_case		=			ReformerModelWithLMHead(__lowercase )

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

			set_model_weights_in_torch(__lowercase       ,					__lowercase       ,					config.hidden_size )

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


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