





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   UpperCamelCase_       (     __a				,  unittest.TestCase    ):
							lowerCamelCase_           =					VideoToVideoSDPipeline
							lowerCamelCase_           =					TEXT_GUIDED_IMAGE_VARIATION_PARAMS.union({"video"}    ) - {"image", "width", "height"}
							lowerCamelCase_           =					TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS.union({"video"}    ) - {"image"}
							lowerCamelCase_           =					PipelineTesterMixin.required_optional_params - {"latents"}
							lowerCamelCase_           =					False

							# No `output_type`.
							lowerCamelCase_           =					frozenset(
							    [
							        "num_inference_steps",
							        "generator",
							        "latents",
							        "return_dict",
							        "callback",
							        "callback_steps",
							    ]    )


							def      _snake_case  (							self    :str		)   ->   Dict:



									"""simple docstring"""



									torch.manual_seed(0		)
									SCREAMING_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							,				)
									SCREAMING_SNAKE_CASE__			=      DDIMScheduler(
									    beta_start=0.0_0_0_8_5							,				beta_end=0.0_1_2							,				beta_schedule="""scaled_linear"""							,				clip_sample=a_							,				set_alpha_to_one=a_							,				)
									torch.manual_seed(0		)
									SCREAMING_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		)
									SCREAMING_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=1000							,				hidden_act="""gelu"""							,				projection_dim=512							,				)
									SCREAMING_SNAKE_CASE__			=      CLIPTextModel(a_		)
									SCREAMING_SNAKE_CASE__			=      CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""		)

									SCREAMING_SNAKE_CASE__			=      {
									    """unet""": unet,
									    """scheduler""": scheduler,
									    """vae""": vae,
									    """text_encoder""": text_encoder,
									    """tokenizer""": tokenizer,
									}
									return components


							def      _snake_case  (							self    :Dict							,				__A    :Tuple							,				__A    :List[Any]=0		)   ->   Optional[int]:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__			=      floats_tensor((1, 3, 3, 32, 32)							,				rng=random.Random(a_		)		).to(a_		)

									if str(a_		).startswith("""mps"""		):
											SCREAMING_SNAKE_CASE__			=      torch.manual_seed(a_		)
									else:
											SCREAMING_SNAKE_CASE__			=      torch.Generator(device=a_		).manual_seed(a_		)
									SCREAMING_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      _snake_case  (							self    :str		)   ->   Any:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__			=      """cpu"""  # ensure determinism for the device-dependent torch.Generator
									SCREAMING_SNAKE_CASE__			=      self.get_dummy_components()
									SCREAMING_SNAKE_CASE__			=      VideoToVideoSDPipeline(**a_		)
									SCREAMING_SNAKE_CASE__			=      sd_pipe.to(a_		)
									sd_pipe.set_progress_bar_config(disable=a_		)

									SCREAMING_SNAKE_CASE__			=      self.get_dummy_inputs(a_		)
									SCREAMING_SNAKE_CASE__			=      """np"""
									SCREAMING_SNAKE_CASE__			=      sd_pipe(**a_		).frames
									SCREAMING_SNAKE_CASE__			=      frames[0][-3:, -3:, -1]

									assert frames[0].shape == (32, 32, 3)
									SCREAMING_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      _snake_case  (							self    :Tuple		)   ->   Tuple:



									"""simple docstring"""



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


							@unittest.skip(reason="""Batching needs to be properly figured out first for this pipeline."""		)
							def      _snake_case  (							self    :Any		)   ->   int:



									"""simple docstring"""



									pass


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



									"""simple docstring"""



									pass


							@unittest.skip(reason="""`num_images_per_prompt` argument is not supported for this pipeline."""		)
							def      _snake_case  (							self    :List[Any]		)   ->   Optional[Any]:



									"""simple docstring"""



									pass


							def      _snake_case  (							self    :int		)   ->   str:



									"""simple docstring"""



									return super().test_progress_bar()




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


							def      _snake_case  (							self    :List[Any]		)   ->   List[str]:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__			=      VideoToVideoSDPipeline.from_pretrained("""cerspense/zeroscope_v2_XL"""							,				torch_dtype=torch.floataa		)
									pipe.enable_model_cpu_offload()

									# 10 frames
									SCREAMING_SNAKE_CASE__			=      torch.Generator(device="""cpu"""		).manual_seed(0		)
									SCREAMING_SNAKE_CASE__			=      torch.randn((1, 10, 3, 1024, 576)							,				generator=a_		)
									SCREAMING_SNAKE_CASE__			=      video.to("""cuda"""		)

									SCREAMING_SNAKE_CASE__			=      """Spiderman is surfing"""

									SCREAMING_SNAKE_CASE__			=      pipe(a_							,				video=a_							,				generator=a_							,				num_inference_steps=3							,				output_type="""pt"""		).frames

									SCREAMING_SNAKE_CASE__			=      np.array([-1.0_4_5_8_9_8_4, -1.1_2_7_9_2_9_7, -0.9_6_6_3_0_8_6, -0.9_1_5_0_3_9_0_6, -0.7_5_0_9_7_6_5_6]		)
									assert np.abs(video_frames.cpu().numpy()[0, 0, 0, 0, -5:] - expected_array		).sum() < 1E-2













import argparse

import tensorflow as tf
import torch

from transformers import BertConfig, BertForMaskedLM
from transformers.models.bert.modeling_bert import (
    BertIntermediate,
    BertLayer,
    BertOutput,
    BertPooler,
    BertSelfAttention,
    BertSelfOutput,
)
from transformers.utils import logging


logging.set_verbosity_info()






def 				SCREAMING_SNAKE_CASE__							(       UpperCamelCase__:							str					,   UpperCamelCase__:							str					,   UpperCamelCase__:							str    ):
		def get_masked_lm_array(UpperCamelCase__:							str    ):
				SCREAMING_SNAKE_CASE__			=      f'''masked_lm/{name}/.ATTRIBUTES/VARIABLE_VALUE'''
				SCREAMING_SNAKE_CASE__			=      tf.train.load_variable(UpperCamelCase__					,   UpperCamelCase__    )

				if "kernel" in name:
						SCREAMING_SNAKE_CASE__			=      array.transpose()

				return torch.from_numpy(UpperCamelCase__    )

		def get_encoder_array(UpperCamelCase__:							str    ):
				SCREAMING_SNAKE_CASE__			=      f'''encoder/{name}/.ATTRIBUTES/VARIABLE_VALUE'''
				SCREAMING_SNAKE_CASE__			=      tf.train.load_variable(UpperCamelCase__					,   UpperCamelCase__    )

				if "kernel" in name:
						SCREAMING_SNAKE_CASE__			=      array.transpose()

				return torch.from_numpy(UpperCamelCase__    )

		def get_encoder_layer_array(UpperCamelCase__:							int					,   UpperCamelCase__:							str    ):
				SCREAMING_SNAKE_CASE__			=      f'''encoder/_transformer_layers/{layer_index}/{name}/.ATTRIBUTES/VARIABLE_VALUE'''
				SCREAMING_SNAKE_CASE__			=      tf.train.load_variable(UpperCamelCase__					,   UpperCamelCase__    )

				if "kernel" in name:
						SCREAMING_SNAKE_CASE__			=      array.transpose()

				return torch.from_numpy(UpperCamelCase__    )

		def get_encoder_attention_layer_array(UpperCamelCase__:							int					,   UpperCamelCase__:							str					,   UpperCamelCase__:							Any    ):
				SCREAMING_SNAKE_CASE__			=      f'''encoder/_transformer_layers/{layer_index}/_attention_layer/{name}/.ATTRIBUTES/VARIABLE_VALUE'''
				SCREAMING_SNAKE_CASE__			=      tf.train.load_variable(UpperCamelCase__					,   UpperCamelCase__    )
				SCREAMING_SNAKE_CASE__			=      array.reshape(UpperCamelCase__    )

				if "kernel" in name:
						SCREAMING_SNAKE_CASE__			=      array.transpose()

				return torch.from_numpy(UpperCamelCase__    )

		print(f'''Loading model based on config from {config_path}...'''    )
		SCREAMING_SNAKE_CASE__			=      BertConfig.from_json_file(UpperCamelCase__    )
		SCREAMING_SNAKE_CASE__			=      BertForMaskedLM(UpperCamelCase__    )

		# Layers
		for layer_index in range(0					,   config.num_hidden_layers    ):
				SCREAMING_SNAKE_CASE__			=      model.bert.encoder.layer[layer_index]

				# Self-attention
				SCREAMING_SNAKE_CASE__			=      layer.attention.self

				SCREAMING_SNAKE_CASE__			=      get_encoder_attention_layer_array(
				    UpperCamelCase__					,   """_query_dense/kernel"""					,   self_attn.query.weight.data.shape    )
				SCREAMING_SNAKE_CASE__			=      get_encoder_attention_layer_array(
				    UpperCamelCase__					,   """_query_dense/bias"""					,   self_attn.query.bias.data.shape    )
				SCREAMING_SNAKE_CASE__			=      get_encoder_attention_layer_array(
				    UpperCamelCase__					,   """_key_dense/kernel"""					,   self_attn.key.weight.data.shape    )
				SCREAMING_SNAKE_CASE__			=      get_encoder_attention_layer_array(
				    UpperCamelCase__					,   """_key_dense/bias"""					,   self_attn.key.bias.data.shape    )
				SCREAMING_SNAKE_CASE__			=      get_encoder_attention_layer_array(
				    UpperCamelCase__					,   """_value_dense/kernel"""					,   self_attn.value.weight.data.shape    )
				SCREAMING_SNAKE_CASE__			=      get_encoder_attention_layer_array(
				    UpperCamelCase__					,   """_value_dense/bias"""					,   self_attn.value.bias.data.shape    )

				# Self-attention Output
				SCREAMING_SNAKE_CASE__			=      layer.attention.output

				SCREAMING_SNAKE_CASE__			=      get_encoder_attention_layer_array(
				    UpperCamelCase__					,   """_output_dense/kernel"""					,   self_output.dense.weight.data.shape    )
				SCREAMING_SNAKE_CASE__			=      get_encoder_attention_layer_array(
				    UpperCamelCase__					,   """_output_dense/bias"""					,   self_output.dense.bias.data.shape    )

				SCREAMING_SNAKE_CASE__			=      get_encoder_layer_array(UpperCamelCase__					,   """_attention_layer_norm/gamma"""    )
				SCREAMING_SNAKE_CASE__			=      get_encoder_layer_array(UpperCamelCase__					,   """_attention_layer_norm/beta"""    )

				# Intermediate
				SCREAMING_SNAKE_CASE__			=      layer.intermediate

				SCREAMING_SNAKE_CASE__			=      get_encoder_layer_array(UpperCamelCase__					,   """_intermediate_dense/kernel"""    )
				SCREAMING_SNAKE_CASE__			=      get_encoder_layer_array(UpperCamelCase__					,   """_intermediate_dense/bias"""    )

				# Output
				SCREAMING_SNAKE_CASE__			=      layer.output

				SCREAMING_SNAKE_CASE__			=      get_encoder_layer_array(UpperCamelCase__					,   """_output_dense/kernel"""    )
				SCREAMING_SNAKE_CASE__			=      get_encoder_layer_array(UpperCamelCase__					,   """_output_dense/bias"""    )

				SCREAMING_SNAKE_CASE__			=      get_encoder_layer_array(UpperCamelCase__					,   """_output_layer_norm/gamma"""    )
				SCREAMING_SNAKE_CASE__			=      get_encoder_layer_array(UpperCamelCase__					,   """_output_layer_norm/beta"""    )

		# Embeddings
		SCREAMING_SNAKE_CASE__			=      get_encoder_array("""_position_embedding_layer/embeddings"""    )
		SCREAMING_SNAKE_CASE__			=      get_encoder_array("""_type_embedding_layer/embeddings"""    )
		SCREAMING_SNAKE_CASE__			=      get_encoder_array("""_embedding_norm_layer/gamma"""    )
		SCREAMING_SNAKE_CASE__			=      get_encoder_array("""_embedding_norm_layer/beta"""    )

		# LM Head
		SCREAMING_SNAKE_CASE__			=      model.cls.predictions.transform

		SCREAMING_SNAKE_CASE__			=      get_masked_lm_array("""dense/kernel"""    )
		SCREAMING_SNAKE_CASE__			=      get_masked_lm_array("""dense/bias"""    )

		SCREAMING_SNAKE_CASE__			=      get_masked_lm_array("""layer_norm/gamma"""    )
		SCREAMING_SNAKE_CASE__			=      get_masked_lm_array("""layer_norm/beta"""    )

		SCREAMING_SNAKE_CASE__			=      get_masked_lm_array("""embedding_table"""    )

		# Pooling
		SCREAMING_SNAKE_CASE__			=      BertPooler(config=UpperCamelCase__    )
		SCREAMING_SNAKE_CASE__			=      get_encoder_array("""_pooler_layer/kernel"""    )
		SCREAMING_SNAKE_CASE__			=      get_encoder_array("""_pooler_layer/bias"""    )

		# Export final model
		model.save_pretrained(UpperCamelCase__    )

		# Integration test - should load without any errors ;)
		SCREAMING_SNAKE_CASE__			=      BertForMaskedLM.from_pretrained(UpperCamelCase__    )
		print(new_model.eval()    )

		print("""Model conversion was done sucessfully!"""    )


if __name__ == "__main__":
	_lowerCamelCase   				=			argparse.ArgumentParser()
	parser.add_argument(
	    '--tf_checkpoint_path', type=str, required=True, help='Path to the TensorFlow Token Dropping checkpoint path.'
	)
	parser.add_argument(
	    '--bert_config_file',
	    type=str,
	    required=True,
	    help='The config json file corresponding to the BERT model. This specifies the model architecture.',
	)
	parser.add_argument(
	    '--pytorch_dump_path',
	    type=str,
	    required=True,
	    help='Path to the output PyTorch model.',
	)
	_lowerCamelCase   				=			parser.parse_args()
	convert_checkpoint_to_pytorch(args.tf_checkpoint_path, args.bert_config_file, args.pytorch_dump_path)