
import gc
import unittest

import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTextModelWithProjection, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    DDPMScheduler,
    PriorTransformer,
    StableUnCLIPPipeline,
    UNetaDConditionModel,
)
from diffusers.pipelines.stable_diffusion.stable_unclip_image_normalizer import StableUnCLIPImageNormalizer
from diffusers.utils.testing_utils import enable_full_determinism, load_numpy, require_torch_gpu, slow, torch_device

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import (
    PipelineKarrasSchedulerTesterMixin,
    PipelineLatentTesterMixin,
    PipelineTesterMixin,
    assert_mean_pixel_difference,
)


enable_full_determinism()







class        lowerCAmelCase__      (   __lowercase		,	__lowercase		,	__lowercase		,	unittest.TestCase	):
						UpperCamelCase_							: Tuple	=		StableUnCLIPPipeline
						UpperCamelCase_							: str	=		TEXT_TO_IMAGE_PARAMS
						UpperCamelCase_							: Any	=		TEXT_TO_IMAGE_BATCH_PARAMS
						UpperCamelCase_							: List[str]	=		TEXT_TO_IMAGE_IMAGE_PARAMS
						UpperCamelCase_							: Optional[int]	=		TEXT_TO_IMAGE_IMAGE_PARAMS

						# TODO(will) Expected attn_bias.stride(1) == 0 to be true, but got false
						UpperCamelCase_							: List[str]	=		False

						def  A_				(			self				)					->      Optional[int]:


								'''simple docstring'''

								_UpperCamelCase						    =	32
								_UpperCamelCase						    =	embedder_hidden_size

								# prior components

								torch.manual_seed(0				)
								_UpperCamelCase						    =	CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""				)

								torch.manual_seed(0				)
								_UpperCamelCase						    =	CLIPTextModelWithProjection(
								    CLIPTextConfig(
								        bos_token_id=0       ,							eos_token_id=2       ,							hidden_size=a       ,							projection_dim=a       ,							intermediate_size=37       ,							layer_norm_eps=1e-05       ,							num_attention_heads=4       ,							num_hidden_layers=5       ,							pad_token_id=1       ,							vocab_size=10_00       ,							)				)

								torch.manual_seed(0				)
								_UpperCamelCase						    =	PriorTransformer(
								    num_attention_heads=2       ,							attention_head_dim=12       ,							embedding_dim=a       ,							num_layers=1       ,							)

								torch.manual_seed(0				)
								_UpperCamelCase						    =	DDPMScheduler(
								    variance_type="""fixed_small_log"""       ,							prediction_type="""sample"""       ,							num_train_timesteps=10_00       ,							clip_sample=a       ,							clip_sample_range=5.0       ,							beta_schedule="""squaredcos_cap_v2"""       ,							)

								# regular denoising components

								torch.manual_seed(0				)
								_UpperCamelCase						    =	StableUnCLIPImageNormalizer(embedding_dim=a				)
								_UpperCamelCase						    =	DDPMScheduler(beta_schedule="""squaredcos_cap_v2"""				)

								torch.manual_seed(0				)
								_UpperCamelCase						    =	CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""				)

								torch.manual_seed(0				)
								_UpperCamelCase						    =	CLIPTextModel(
								    CLIPTextConfig(
								        bos_token_id=0       ,							eos_token_id=2       ,							hidden_size=a       ,							projection_dim=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       ,							)				)

								torch.manual_seed(0				)
								_UpperCamelCase						    =	UNetaDConditionModel(
								    sample_size=32       ,							in_channels=4       ,							out_channels=4       ,							down_block_types=("""CrossAttnDownBlock2D""", """DownBlock2D""")       ,							up_block_types=("""UpBlock2D""", """CrossAttnUpBlock2D""")       ,							block_out_channels=(32, 64)       ,							attention_head_dim=(2, 4)       ,							class_embed_type="""projection"""       ,							projection_class_embeddings_input_dim=embedder_projection_dim * 2       ,							cross_attention_dim=a       ,							layers_per_block=1       ,							upcast_attention=a       ,							use_linear_projection=a       ,							)

								torch.manual_seed(0				)
								_UpperCamelCase						    =	DDIMScheduler(
								    beta_schedule="""scaled_linear"""       ,							beta_start=0.0_0085       ,							beta_end=0.012       ,							prediction_type="""v_prediction"""       ,							set_alpha_to_one=a       ,							steps_offset=1       ,							)

								torch.manual_seed(0				)
								_UpperCamelCase						    =	AutoencoderKL()

								_UpperCamelCase						    =	{
								    # prior components
								    """prior_tokenizer""": prior_tokenizer,
								    """prior_text_encoder""": prior_text_encoder,
								    """prior""": prior,
								    """prior_scheduler""": prior_scheduler,
								    # image noising components
								    """image_normalizer""": image_normalizer,
								    """image_noising_scheduler""": image_noising_scheduler,
								    # regular denoising components
								    """tokenizer""": tokenizer,
								    """text_encoder""": text_encoder,
								    """unet""": unet,
								    """scheduler""": scheduler,
								    """vae""": vae,
								}

								return components

						def  A_				(			self       ,							a       ,							a=0				)					->      Union[str, Any]:


								'''simple docstring'''

								if str(a				).startswith("""mps"""				):
										_UpperCamelCase						    =	torch.manual_seed(a				)
								else:
										_UpperCamelCase						    =	torch.Generator(device=a				).manual_seed(a				)
								_UpperCamelCase						    =	{
								    """prompt""": """A painting of a squirrel eating a burger""",
								    """generator""": generator,
								    """num_inference_steps""": 2,
								    """prior_num_inference_steps""": 2,
								    """output_type""": """numpy""",
								}
								return inputs

						def  A_				(			self				)					->      Union[str, Any]:


								'''simple docstring'''

								_UpperCamelCase						    =	torch_device == """cpu"""

								self._test_attention_slicing_forward_pass(test_max_difference=a				)







						def  A_				(			self				)					->      Any:


								'''simple docstring'''

								_UpperCamelCase						    =	torch_device in ["""cpu""", """mps"""]

								self._test_inference_batch_single_identical(test_max_difference=a				)







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

						def  A_				(			self				)					->      Union[str, Any]:


								'''simple docstring'''

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

						def  A_				(			self				)					->      Optional[Any]:


								'''simple docstring'''

								_UpperCamelCase						    =	load_numpy(
								    """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/stable_unclip/stable_unclip_2_1_l_anime_turtle_fp16.npy"""				)

								_UpperCamelCase						    =	StableUnCLIPPipeline.from_pretrained("""fusing/stable-unclip-2-1-l"""       ,							torch_dtype=torch.floataa				)
								pipe.to(a				)
								pipe.set_progress_bar_config(disable=a				)
								# stable unclip will oom when integration tests are run on a V100,
								# so turn on memory savings
								pipe.enable_attention_slicing()
								pipe.enable_sequential_cpu_offload()

								_UpperCamelCase						    =	torch.Generator(device="""cpu"""				).manual_seed(0				)
								_UpperCamelCase						    =	pipe("""anime turle"""       ,							generator=a       ,							output_type="""np"""				)

								_UpperCamelCase						    =	output.images[0]

								assert image.shape == (7_68, 7_68, 3)

								assert_mean_pixel_difference(a       ,							a				)







						def  A_				(			self				)					->      Optional[Any]:


								'''simple docstring'''

								torch.cuda.empty_cache()
								torch.cuda.reset_max_memory_allocated()
								torch.cuda.reset_peak_memory_stats()

								_UpperCamelCase						    =	StableUnCLIPPipeline.from_pretrained("""fusing/stable-unclip-2-1-l"""       ,							torch_dtype=torch.floataa				)
								_UpperCamelCase						    =	pipe.to(a				)
								pipe.set_progress_bar_config(disable=a				)
								pipe.enable_attention_slicing()
								pipe.enable_sequential_cpu_offload()

								_UpperCamelCase						    =	pipe(
								    """anime turtle"""       ,							prior_num_inference_steps=2       ,							num_inference_steps=2       ,							output_type="""np"""       ,							)

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





import argparse
import os
import re

import tensorflow as tf
import torch

from transformers import BertConfig, BertModel
from transformers.utils import logging


logging.set_verbosity_info()
lowerCamelCase__	=							logging.get_logger(__name__)


def 							__A(lowerCAmelCase						,  lowerCAmelCase						,  lowerCAmelCase )		->				Any:

		"""simple docstring"""






		_UpperCamelCase						    =	os.path.abspath(lowerCAmelCase )
		logger.info(F'Converting TensorFlow checkpoint from {tf_path}' )
		# Load weights from TF model
		_UpperCamelCase						    =	tf.train.list_variables(lowerCAmelCase )
		_UpperCamelCase						    =	[]
		_UpperCamelCase						    =	[]
		_UpperCamelCase						    =	[]
		for full_name, shape in init_vars:
				# logger.info(f"Loading TF weight {name} with shape {shape}")
				_UpperCamelCase						    =	full_name.split("""/""" )
				if full_name == "_CHECKPOINTABLE_OBJECT_GRAPH" or name[0] in ["global_step", "save_counter"]:
						logger.info(F'Skipping non-model layer {full_name}' )
						continue
				if "optimizer" in full_name:
						logger.info(F'Skipping optimization layer {full_name}' )
						continue
				if name[0] == "model":
						# ignore initial 'model'
						_UpperCamelCase						    =	name[1:]
				# figure out how many levels deep the name is
				_UpperCamelCase						    =	0
				for _name in name:
						if _name.startswith("""layer_with_weights""" ):
								depth += 1
						else:
								break
				layer_depth.append(lowerCAmelCase )
				# read data
				_UpperCamelCase						    =	tf.train.load_variable(lowerCAmelCase						,  lowerCAmelCase )
				names.append("""/""".join(lowerCAmelCase ) )
				arrays.append(lowerCAmelCase )
		logger.info(F'Read a total of {len(lowerCAmelCase ):,} layers' )

		# Sanity check
		if len(set(lowerCAmelCase ) ) != 1:
				raise ValueError(F'Found layer names with different depths (layer depth {list(set(lowerCAmelCase ) )})' )
		_UpperCamelCase						    =	list(set(lowerCAmelCase ) )[0]
		if layer_depth != 1:
				raise ValueError(
				    """The model contains more than just the embedding/encoder layers. This script does not handle MLM/NSP"""
				    """ heads.""" )

		# convert layers
		logger.info("""Converting weights...""" )
		for full_name, array in zip(lowerCAmelCase						,  lowerCAmelCase ):
				_UpperCamelCase						    =	full_name.split("""/""" )
				_UpperCamelCase						    =	model
				_UpperCamelCase						    =	[]
				for i, m_name in enumerate(lowerCAmelCase ):
						if m_name == ".ATTRIBUTES":
								# variable names end with .ATTRIBUTES/VARIABLE_VALUE
								break
						if m_name.startswith("""layer_with_weights""" ):
								_UpperCamelCase						    =	int(m_name.split("""-""" )[-1] )
								if layer_num <= 2:
										# embedding layers
										# layer_num 0: word_embeddings
										# layer_num 1: position_embeddings
										# layer_num 2: token_type_embeddings
										continue
								elif layer_num == 3:
										# embedding LayerNorm
										trace.extend(["""embeddings""", """LayerNorm"""] )
										_UpperCamelCase						    =	getattr(lowerCAmelCase						,  """embeddings""" )
										_UpperCamelCase						    =	getattr(lowerCAmelCase						,  """LayerNorm""" )
								elif layer_num > 3 and layer_num < config.num_hidden_layers + 4:
										# encoder layers
										trace.extend(["""encoder""", """layer""", str(layer_num - 4 )] )
										_UpperCamelCase						    =	getattr(lowerCAmelCase						,  """encoder""" )
										_UpperCamelCase						    =	getattr(lowerCAmelCase						,  """layer""" )
										_UpperCamelCase						    =	pointer[layer_num - 4]
								elif layer_num == config.num_hidden_layers + 4:
										# pooler layer
										trace.extend(["""pooler""", """dense"""] )
										_UpperCamelCase						    =	getattr(lowerCAmelCase						,  """pooler""" )
										_UpperCamelCase						    =	getattr(lowerCAmelCase						,  """dense""" )
						elif m_name == "embeddings":
								trace.append("""embeddings""" )
								_UpperCamelCase						    =	getattr(lowerCAmelCase						,  """embeddings""" )
								if layer_num == 0:
										trace.append("""word_embeddings""" )
										_UpperCamelCase						    =	getattr(lowerCAmelCase						,  """word_embeddings""" )
								elif layer_num == 1:
										trace.append("""position_embeddings""" )
										_UpperCamelCase						    =	getattr(lowerCAmelCase						,  """position_embeddings""" )
								elif layer_num == 2:
										trace.append("""token_type_embeddings""" )
										_UpperCamelCase						    =	getattr(lowerCAmelCase						,  """token_type_embeddings""" )
								else:
										raise ValueError(F'Unknown embedding layer with name {full_name}' )
								trace.append("""weight""" )
								_UpperCamelCase						    =	getattr(lowerCAmelCase						,  """weight""" )
						elif m_name == "_attention_layer":
								# self-attention layer
								trace.extend(["""attention""", """self"""] )
								_UpperCamelCase						    =	getattr(lowerCAmelCase						,  """attention""" )
								_UpperCamelCase						    =	getattr(lowerCAmelCase						,  """self""" )
						elif m_name == "_attention_layer_norm":
								# output attention norm
								trace.extend(["""attention""", """output""", """LayerNorm"""] )
								_UpperCamelCase						    =	getattr(lowerCAmelCase						,  """attention""" )
								_UpperCamelCase						    =	getattr(lowerCAmelCase						,  """output""" )
								_UpperCamelCase						    =	getattr(lowerCAmelCase						,  """LayerNorm""" )
						elif m_name == "_attention_output_dense":
								# output attention dense
								trace.extend(["""attention""", """output""", """dense"""] )
								_UpperCamelCase						    =	getattr(lowerCAmelCase						,  """attention""" )
								_UpperCamelCase						    =	getattr(lowerCAmelCase						,  """output""" )
								_UpperCamelCase						    =	getattr(lowerCAmelCase						,  """dense""" )
						elif m_name == "_output_dense":
								# output dense
								trace.extend(["""output""", """dense"""] )
								_UpperCamelCase						    =	getattr(lowerCAmelCase						,  """output""" )
								_UpperCamelCase						    =	getattr(lowerCAmelCase						,  """dense""" )
						elif m_name == "_output_layer_norm":
								# output dense
								trace.extend(["""output""", """LayerNorm"""] )
								_UpperCamelCase						    =	getattr(lowerCAmelCase						,  """output""" )
								_UpperCamelCase						    =	getattr(lowerCAmelCase						,  """LayerNorm""" )
						elif m_name == "_key_dense":
								# attention key
								trace.append("""key""" )
								_UpperCamelCase						    =	getattr(lowerCAmelCase						,  """key""" )
						elif m_name == "_query_dense":
								# attention query
								trace.append("""query""" )
								_UpperCamelCase						    =	getattr(lowerCAmelCase						,  """query""" )
						elif m_name == "_value_dense":
								# attention value
								trace.append("""value""" )
								_UpperCamelCase						    =	getattr(lowerCAmelCase						,  """value""" )
						elif m_name == "_intermediate_dense":
								# attention intermediate dense
								trace.extend(["""intermediate""", """dense"""] )
								_UpperCamelCase						    =	getattr(lowerCAmelCase						,  """intermediate""" )
								_UpperCamelCase						    =	getattr(lowerCAmelCase						,  """dense""" )
						elif m_name == "_output_layer_norm":
								# output layer norm
								trace.append("""output""" )
								_UpperCamelCase						    =	getattr(lowerCAmelCase						,  """output""" )
						# weights & biases
						elif m_name in ["bias", "beta"]:
								trace.append("""bias""" )
								_UpperCamelCase						    =	getattr(lowerCAmelCase						,  """bias""" )
						elif m_name in ["kernel", "gamma"]:
								trace.append("""weight""" )
								_UpperCamelCase						    =	getattr(lowerCAmelCase						,  """weight""" )
						else:
								logger.warning(F'Ignored {m_name}' )
        # for certain layers reshape is necessary
				_UpperCamelCase						    =	""".""".join(lowerCAmelCase )
				if re.match(R"""(\S+)\.attention\.self\.(key|value|query)\.(bias|weight)"""						,  lowerCAmelCase ) or re.match(
				    R"""(\S+)\.attention\.output\.dense\.weight"""						,  lowerCAmelCase ):
						_UpperCamelCase						    =	array.reshape(pointer.data.shape )
				if "kernel" in full_name:
						_UpperCamelCase						    =	array.transpose()
				if pointer.shape == array.shape:
						_UpperCamelCase						    =	torch.from_numpy(lowerCAmelCase )
				else:
						raise ValueError(
						    F'Shape mismatch in layer {full_name}: Model expects shape {pointer.shape} but layer contains shape:'
						    F' {array.shape}' )
				logger.info(F'Successfully set variable {full_name} to PyTorch layer {trace}' )
		return model


def 							__A(lowerCAmelCase						,  lowerCAmelCase						,  lowerCAmelCase )		->				Tuple:

		"""simple docstring"""






		logger.info(F'Loading model based on config from {config_path}...' )
		_UpperCamelCase						    =	BertConfig.from_json_file(lowerCAmelCase )
		_UpperCamelCase						    =	BertModel(lowerCAmelCase )

		# Load weights from checkpoint
		logger.info(F'Loading weights from checkpoint {tf_checkpoint_path}...' )
		load_tfa_weights_in_bert(lowerCAmelCase						,  lowerCAmelCase						,  lowerCAmelCase )

		# Save pytorch-model
		logger.info(F'Saving PyTorch model to {pytorch_dump_path}...' )
		torch.save(model.state_dict()						,  lowerCAmelCase )


if __name__ == "__main__":
		lowerCamelCase__	=							argparse.ArgumentParser()
		parser.add_argument(
		    "--tf_checkpoint_path", type=str, required=True, help="Path to the TensorFlow 2.x 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 (must include filename).",
		)
		lowerCamelCase__	=							parser.parse_args()
		convert_tfa_checkpoint_to_pytorch(args.tf_checkpoint_path, args.bert_config_file, args.pytorch_dump_path)



