


'''simple docstring'''






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()
__UpperCAmelCase       							=				logging.get_logger(__name__)


def    _snake_case  (	A					,       A					,       A			)		->       Union[str, Any]:
							lowerCAmelCase__			  =      os.path.abspath(A			)
							logger.info(F"""Converting TensorFlow checkpoint from {tf_path}"""			)
							# Load weights from TF model
							lowerCAmelCase__			  =      tf.train.list_variables(A			)
							lowerCAmelCase__			  =      []
							lowerCAmelCase__			  =      []
							lowerCAmelCase__			  =      []
							for full_name, shape in init_vars:
														# logger.info(f"Loading TF weight {name} with shape {shape}")
														lowerCAmelCase__			  =      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'
																					lowerCAmelCase__			  =      name[1:]
														# figure out how many levels deep the name is
														lowerCAmelCase__			  =      0
														for _name in name:
																					if _name.startswith('''layer_with_weights'''			):
																												depth += 1
																					else:
																												break
														layer_depth.append(A			)
														# read data
														lowerCAmelCase__			  =      tf.train.load_variable(A					,       A			)
														names.append('''/'''.join(A			)			)
														arrays.append(A			)
							logger.info(F"""Read a total of {len(A			):,} layers"""			)

							# Sanity check
							if len(set(A			)			) != 1:
														raise ValueError(F"""Found layer names with different depths (layer depth {list(set(A			)			)})"""			)
							lowerCAmelCase__			  =      list(set(A			)			)[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(A					,       A			):
														lowerCAmelCase__			  =      full_name.split('''/'''			)
														lowerCAmelCase__			  =      model
														lowerCAmelCase__			  =      []
														for i, m_name in enumerate(A			):
																					if m_name == ".ATTRIBUTES":
																												# variable names end with .ATTRIBUTES/VARIABLE_VALUE
																												break
																					if m_name.startswith('''layer_with_weights'''			):
																												lowerCAmelCase__			  =      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''']			)
																																			lowerCAmelCase__			  =      getattr(A					,       '''embeddings'''			)
																																			lowerCAmelCase__			  =      getattr(A					,       '''LayerNorm'''			)
																												elif layer_num > 3 and layer_num < config.num_hidden_layers + 4:
																																			# encoder layers
																																			trace.extend(['''encoder''', '''layer''', str(layer_num - 4			)]			)
																																			lowerCAmelCase__			  =      getattr(A					,       '''encoder'''			)
																																			lowerCAmelCase__			  =      getattr(A					,       '''layer'''			)
																																			lowerCAmelCase__			  =      pointer[layer_num - 4]
																												elif layer_num == config.num_hidden_layers + 4:
																																			# pooler layer
																																			trace.extend(['''pooler''', '''dense''']			)
																																			lowerCAmelCase__			  =      getattr(A					,       '''pooler'''			)
																																			lowerCAmelCase__			  =      getattr(A					,       '''dense'''			)
																					elif m_name == "embeddings":
																												trace.append('''embeddings'''			)
																												lowerCAmelCase__			  =      getattr(A					,       '''embeddings'''			)
																												if layer_num == 0:
																																			trace.append('''word_embeddings'''			)
																																			lowerCAmelCase__			  =      getattr(A					,       '''word_embeddings'''			)
																												elif layer_num == 1:
																																			trace.append('''position_embeddings'''			)
																																			lowerCAmelCase__			  =      getattr(A					,       '''position_embeddings'''			)
																												elif layer_num == 2:
																																			trace.append('''token_type_embeddings'''			)
																																			lowerCAmelCase__			  =      getattr(A					,       '''token_type_embeddings'''			)
																												else:
																																			raise ValueError(F"""Unknown embedding layer with name {full_name}"""			)
																												trace.append('''weight'''			)
																												lowerCAmelCase__			  =      getattr(A					,       '''weight'''			)
																					elif m_name == "_attention_layer":
																												# self-attention layer
																												trace.extend(['''attention''', '''self''']			)
																												lowerCAmelCase__			  =      getattr(A					,       '''attention'''			)
																												lowerCAmelCase__			  =      getattr(A					,       '''self'''			)
																					elif m_name == "_attention_layer_norm":
																												# output attention norm
																												trace.extend(['''attention''', '''output''', '''LayerNorm''']			)
																												lowerCAmelCase__			  =      getattr(A					,       '''attention'''			)
																												lowerCAmelCase__			  =      getattr(A					,       '''output'''			)
																												lowerCAmelCase__			  =      getattr(A					,       '''LayerNorm'''			)
																					elif m_name == "_attention_output_dense":
																												# output attention dense
																												trace.extend(['''attention''', '''output''', '''dense''']			)
																												lowerCAmelCase__			  =      getattr(A					,       '''attention'''			)
																												lowerCAmelCase__			  =      getattr(A					,       '''output'''			)
																												lowerCAmelCase__			  =      getattr(A					,       '''dense'''			)
																					elif m_name == "_output_dense":
																												# output dense
																												trace.extend(['''output''', '''dense''']			)
																												lowerCAmelCase__			  =      getattr(A					,       '''output'''			)
																												lowerCAmelCase__			  =      getattr(A					,       '''dense'''			)
																					elif m_name == "_output_layer_norm":
																												# output dense
																												trace.extend(['''output''', '''LayerNorm''']			)
																												lowerCAmelCase__			  =      getattr(A					,       '''output'''			)
																												lowerCAmelCase__			  =      getattr(A					,       '''LayerNorm'''			)
																					elif m_name == "_key_dense":
																												# attention key
																												trace.append('''key'''			)
																												lowerCAmelCase__			  =      getattr(A					,       '''key'''			)
																					elif m_name == "_query_dense":
																												# attention query
																												trace.append('''query'''			)
																												lowerCAmelCase__			  =      getattr(A					,       '''query'''			)
																					elif m_name == "_value_dense":
																												# attention value
																												trace.append('''value'''			)
																												lowerCAmelCase__			  =      getattr(A					,       '''value'''			)
																					elif m_name == "_intermediate_dense":
																												# attention intermediate dense
																												trace.extend(['''intermediate''', '''dense''']			)
																												lowerCAmelCase__			  =      getattr(A					,       '''intermediate'''			)
																												lowerCAmelCase__			  =      getattr(A					,       '''dense'''			)
																					elif m_name == "_output_layer_norm":
																												# output layer norm
																												trace.append('''output'''			)
																												lowerCAmelCase__			  =      getattr(A					,       '''output'''			)
																					# weights & biases
																					elif m_name in ["bias", "beta"]:
																												trace.append('''bias'''			)
																												lowerCAmelCase__			  =      getattr(A					,       '''bias'''			)
																					elif m_name in ["kernel", "gamma"]:
																												trace.append('''weight'''			)
																												lowerCAmelCase__			  =      getattr(A					,       '''weight'''			)
																					else:
																												logger.warning(F"""Ignored {m_name}"""			)
        # for certain layers reshape is necessary
														lowerCAmelCase__			  =      '''.'''.join(A			)
														if re.match(R'''(\S+)\.attention\.self\.(key|value|query)\.(bias|weight)'''					,       A			) or re.match(
														    R'''(\S+)\.attention\.output\.dense\.weight'''					,       A			):
																					lowerCAmelCase__			  =      array.reshape(pointer.data.shape			)
														if "kernel" in full_name:
																					lowerCAmelCase__			  =      array.transpose()
														if pointer.shape == array.shape:
																					lowerCAmelCase__			  =      torch.from_numpy(A			)
														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    _snake_case  (	A					,       A					,       A			)		->       Any:
							# Instantiate model
							logger.info(F"""Loading model based on config from {config_path}..."""			)
							lowerCAmelCase__			  =      BertConfig.from_json_file(A			)
							lowerCAmelCase__			  =      BertModel(A			)

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

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


if __name__ == "__main__":
		__UpperCAmelCase       							=				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).''',
		)
		__UpperCAmelCase       							=				parser.parse_args()
		convert_tfa_checkpoint_to_pytorch(args.tf_checkpoint_path, args.bert_config_file, args.pytorch_dump_path)



'''simple docstring'''






import gc
import unittest

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

from diffusers import AltDiffusionPipeline, AutoencoderKL, DDIMScheduler, PNDMScheduler, UNetaDConditionModel
from diffusers.pipelines.alt_diffusion.modeling_roberta_series import (
    RobertaSeriesConfig,
    RobertaSeriesModelWithTransformation,
)
from diffusers.utils import 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
from ..test_pipelines_common import PipelineKarrasSchedulerTesterMixin, PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()




class 						a__							(				a__							,							a__							,							a__							,							unittest.TestCase    ):





							'''simple docstring'''






							lowercase__     :  Optional[Any]      		= AltDiffusionPipeline
							lowercase__     :  Dict      		= TEXT_TO_IMAGE_PARAMS
							lowercase__     :  str      		= TEXT_TO_IMAGE_BATCH_PARAMS
							lowercase__     :  Any      		= TEXT_TO_IMAGE_IMAGE_PARAMS
							lowercase__     :  List[str]      		= TEXT_TO_IMAGE_IMAGE_PARAMS







							def     __SCREAMING_SNAKE_CASE				(						self  )			->				str:
														torch.manual_seed(0  )
														lowerCAmelCase__			  =      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   ,       )
														lowerCAmelCase__			  =      DDIMScheduler(
														    beta_start=0.00_085   ,       beta_end=0.012   ,       beta_schedule='''scaled_linear'''   ,       clip_sample=lowerCamelCase_   ,       set_alpha_to_one=lowerCamelCase_   ,       )
														torch.manual_seed(0  )
														lowerCAmelCase__			  =      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   ,       )

														# TODO: address the non-deterministic text encoder (fails for save-load tests)
														# torch.manual_seed(0)
														# text_encoder_config = RobertaSeriesConfig(
														#     hidden_size=32,
														#     project_dim=32,
														#     intermediate_size=37,
														#     layer_norm_eps=1e-05,
														#     num_attention_heads=4,
														#     num_hidden_layers=5,
														#     vocab_size=5002,
														# )
														# text_encoder = RobertaSeriesModelWithTransformation(text_encoder_config)

														torch.manual_seed(0  )
														lowerCAmelCase__			  =      CLIPTextConfig(
														    bos_token_id=0   ,       eos_token_id=2   ,       hidden_size=32   ,       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=50_02   ,       )
														lowerCAmelCase__			  =      CLIPTextModel(lowerCamelCase_  )

														lowerCAmelCase__			  =      XLMRobertaTokenizer.from_pretrained('''hf-internal-testing/tiny-xlm-roberta'''  )
														lowerCAmelCase__			  =      77

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







							def     __SCREAMING_SNAKE_CASE				(						self   ,       lowerCamelCase_   ,       lowerCamelCase_=0  )			->				List[str]:
														if str(lowerCamelCase_  ).startswith('''mps'''  ):
																					lowerCAmelCase__			  =      torch.manual_seed(lowerCamelCase_  )
														else:
																					lowerCAmelCase__			  =      torch.Generator(device=lowerCamelCase_  ).manual_seed(lowerCamelCase_  )
														lowerCAmelCase__			  =      {
														    '''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     __SCREAMING_SNAKE_CASE				(						self  )			->				Union[str, Any]:
														super().test_attention_slicing_forward_pass(expected_max_diff=3e-3  )







							def     __SCREAMING_SNAKE_CASE				(						self  )			->				Optional[int]:
														super().test_inference_batch_single_identical(expected_max_diff=3e-3  )







							def     __SCREAMING_SNAKE_CASE				(						self  )			->				List[str]:
														lowerCAmelCase__			  =      '''cpu'''  # ensure determinism for the device-dependent torch.Generator

														lowerCAmelCase__			  =      self.get_dummy_components()
														torch.manual_seed(0  )
														lowerCAmelCase__			  =      RobertaSeriesConfig(
														    hidden_size=32   ,       project_dim=32   ,       intermediate_size=37   ,       layer_norm_eps=1e-05   ,       num_attention_heads=4   ,       num_hidden_layers=5   ,       vocab_size=50_02   ,       )
														# TODO: remove after fixing the non-deterministic text encoder
														lowerCAmelCase__			  =      RobertaSeriesModelWithTransformation(lowerCamelCase_  )
														lowerCAmelCase__			  =      text_encoder

														lowerCAmelCase__			  =      AltDiffusionPipeline(**lowerCamelCase_  )
														lowerCAmelCase__			  =      alt_pipe.to(lowerCamelCase_  )
														alt_pipe.set_progress_bar_config(disable=lowerCamelCase_  )

														lowerCAmelCase__			  =      self.get_dummy_inputs(lowerCamelCase_  )
														lowerCAmelCase__			  =      '''A photo of an astronaut'''
														lowerCAmelCase__			  =      alt_pipe(**lowerCamelCase_  )
														lowerCAmelCase__			  =      output.images
														lowerCAmelCase__			  =      image[0, -3:, -3:, -1]

														assert image.shape == (1, 64, 64, 3)
														lowerCAmelCase__			  =      np.array(
														    [0.5_748_162, 0.60_447_145, 0.48_821_217, 0.50_100_636, 0.5_431_185, 0.45_763_683, 0.49_657_696, 0.48_132_733, 0.47_573_093]  )

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







							def     __SCREAMING_SNAKE_CASE				(						self  )			->				List[Any]:
														lowerCAmelCase__			  =      '''cpu'''  # ensure determinism for the device-dependent torch.Generator

														lowerCAmelCase__			  =      self.get_dummy_components()
														lowerCAmelCase__			  =      PNDMScheduler(skip_prk_steps=lowerCamelCase_  )
														torch.manual_seed(0  )
														lowerCAmelCase__			  =      RobertaSeriesConfig(
														    hidden_size=32   ,       project_dim=32   ,       intermediate_size=37   ,       layer_norm_eps=1e-05   ,       num_attention_heads=4   ,       num_hidden_layers=5   ,       vocab_size=50_02   ,       )
														# TODO: remove after fixing the non-deterministic text encoder
														lowerCAmelCase__			  =      RobertaSeriesModelWithTransformation(lowerCamelCase_  )
														lowerCAmelCase__			  =      text_encoder
														lowerCAmelCase__			  =      AltDiffusionPipeline(**lowerCamelCase_  )
														lowerCAmelCase__			  =      alt_pipe.to(lowerCamelCase_  )
														alt_pipe.set_progress_bar_config(disable=lowerCamelCase_  )

														lowerCAmelCase__			  =      self.get_dummy_inputs(lowerCamelCase_  )
														lowerCAmelCase__			  =      alt_pipe(**lowerCamelCase_  )
														lowerCAmelCase__			  =      output.images
														lowerCAmelCase__			  =      image[0, -3:, -3:, -1]

														assert image.shape == (1, 64, 64, 3)
														lowerCAmelCase__			  =      np.array(
														    [0.51_605_093, 0.5_707_241, 0.47_365_507, 0.50_578_886, 0.5_633_877, 0.4_642_503, 0.5_182_081, 0.48_763_484, 0.49_084_237]  )

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




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





							'''simple docstring'''






							def     __SCREAMING_SNAKE_CASE				(						self  )			->				Tuple:
														# clean up the VRAM after each test
														super().tearDown()
														gc.collect()
														torch.cuda.empty_cache()







							def     __SCREAMING_SNAKE_CASE				(						self  )			->				List[Any]:
														# make sure here that pndm scheduler skips prk
														lowerCAmelCase__			  =      AltDiffusionPipeline.from_pretrained('''BAAI/AltDiffusion'''   ,       safety_checker=lowerCamelCase_  )
														lowerCAmelCase__			  =      alt_pipe.to(lowerCamelCase_  )
														alt_pipe.set_progress_bar_config(disable=lowerCamelCase_  )

														lowerCAmelCase__			  =      '''A painting of a squirrel eating a burger'''
														lowerCAmelCase__			  =      torch.manual_seed(0  )
														lowerCAmelCase__			  =      alt_pipe([prompt]   ,       generator=lowerCamelCase_   ,       guidance_scale=6.0   ,       num_inference_steps=20   ,       output_type='''np'''  )

														lowerCAmelCase__			  =      output.images

														lowerCAmelCase__			  =      image[0, -3:, -3:, -1]

														assert image.shape == (1, 5_12, 5_12, 3)
														lowerCAmelCase__			  =      np.array([0.1_010, 0.0_800, 0.0_794, 0.0_885, 0.0_843, 0.0_762, 0.0_769, 0.0_729, 0.0_586]  )

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







							def     __SCREAMING_SNAKE_CASE				(						self  )			->				Any:
														lowerCAmelCase__			  =      DDIMScheduler.from_pretrained('''BAAI/AltDiffusion'''   ,       subfolder='''scheduler'''  )

														lowerCAmelCase__			  =      AltDiffusionPipeline.from_pretrained('''BAAI/AltDiffusion'''   ,       scheduler=lowerCamelCase_   ,       safety_checker=lowerCamelCase_  )
														lowerCAmelCase__			  =      alt_pipe.to(lowerCamelCase_  )
														alt_pipe.set_progress_bar_config(disable=lowerCamelCase_  )

														lowerCAmelCase__			  =      '''A painting of a squirrel eating a burger'''
														lowerCAmelCase__			  =      torch.manual_seed(0  )

														lowerCAmelCase__			  =      alt_pipe([prompt]   ,       generator=lowerCamelCase_   ,       num_inference_steps=2   ,       output_type='''numpy'''  )
														lowerCAmelCase__			  =      output.images

														lowerCAmelCase__			  =      image[0, -3:, -3:, -1]

														assert image.shape == (1, 5_12, 5_12, 3)
														lowerCAmelCase__			  =      np.array([0.4_019, 0.4_052, 0.3_810, 0.4_119, 0.3_916, 0.3_982, 0.4_651, 0.4_195, 0.5_323]  )

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