



"""simple docstring"""





import gc
import unittest

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

from diffusers import TransformeraDModel, VQDiffusionPipeline, VQDiffusionScheduler, VQModel
from diffusers.pipelines.vq_diffusion.pipeline_vq_diffusion import LearnedClassifierFreeSamplingEmbeddings
from diffusers.utils import load_numpy, slow, torch_device
from diffusers.utils.testing_utils import require_torch_gpu


__lowerCAmelCase       :					Optional[Any]      					=					False



class    _lowerCAmelCase    (				unittest.TestCase		):





		"""simple docstring"""






		def    UpperCAmelCase__							(					self   ) ->						Any:







								'''simple docstring'''

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





		@property
		def    UpperCAmelCase__							(					self   ) ->						Any:







								'''simple docstring'''

								return 1_2





		@property
		def    UpperCAmelCase__							(					self   ) ->						Dict:







								'''simple docstring'''

								return 1_2





		@property
		def    UpperCAmelCase__							(					self   ) ->						List[Any]:







								'''simple docstring'''

								return 3_2





		@property
		def    UpperCAmelCase__							(					self   ) ->						Optional[Any]:







								'''simple docstring'''

								torch.manual_seed(0   )
								snake_case_   : List[Any]						=			VQModel(
								    block_out_channels=[3_2, 6_4]							,    in_channels=3							,    out_channels=3							,    down_block_types=["""DownEncoderBlock2D""", """DownEncoderBlock2D"""]							,    up_block_types=["""UpDecoderBlock2D""", """UpDecoderBlock2D"""]							,    latent_channels=3							,    num_vq_embeddings=self.num_embed							,    vq_embed_dim=3							,    )
								return model





		@property
		def    UpperCAmelCase__							(					self   ) ->						List[Any]:







								'''simple docstring'''

								snake_case_   : int						=			CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""   )
								return tokenizer





		@property
		def    UpperCAmelCase__							(					self   ) ->						int:







								'''simple docstring'''

								torch.manual_seed(0   )
								snake_case_   : Dict						=			CLIPTextConfig(
								    bos_token_id=0							,    eos_token_id=2							,    hidden_size=self.text_embedder_hidden_size							,    intermediate_size=3_7							,    layer_norm_eps=1E-05							,    num_attention_heads=4							,    num_hidden_layers=5							,    pad_token_id=1							,    vocab_size=1_0_0_0							,    )
								return CLIPTextModel(_lowercase   )





		@property
		def    UpperCAmelCase__							(					self   ) ->						str:







								'''simple docstring'''

								torch.manual_seed(0   )

								snake_case_   : Union[str, Any]						=			1_2
								snake_case_   : Tuple						=			1_2

								snake_case_   : Tuple						=			{
								    """attention_bias""": True,
								    """cross_attention_dim""": 3_2,
								    """attention_head_dim""": height * width,
								    """num_attention_heads""": 1,
								    """num_vector_embeds""": self.num_embed,
								    """num_embeds_ada_norm""": self.num_embeds_ada_norm,
								    """norm_num_groups""": 3_2,
								    """sample_size""": width,
								    """activation_fn""": """geglu-approximate""",
								}

								snake_case_   : Optional[Any]						=			TransformeraDModel(**_lowercase   )
								return model





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







								'''simple docstring'''

								snake_case_   : str						=			"""cpu"""

								snake_case_   : List[str]						=			self.dummy_vqvae
								snake_case_   : Any						=			self.dummy_text_encoder
								snake_case_   : Tuple						=			self.dummy_tokenizer
								snake_case_   : int						=			self.dummy_transformer
								snake_case_   : int						=			VQDiffusionScheduler(self.num_embed   )
								snake_case_   : Dict						=			LearnedClassifierFreeSamplingEmbeddings(learnable=_lowercase   )

								snake_case_   : Optional[Any]						=			VQDiffusionPipeline(
								    vqvae=_lowercase							,    text_encoder=_lowercase							,    tokenizer=_lowercase							,    transformer=_lowercase							,    scheduler=_lowercase							,    learned_classifier_free_sampling_embeddings=_lowercase							,    )
								snake_case_   : int						=			pipe.to(_lowercase   )
								pipe.set_progress_bar_config(disable=_lowercase   )

								snake_case_   : List[Any]						=			"""teddy bear playing in the pool"""

								snake_case_   : Dict						=			torch.Generator(device=_lowercase   ).manual_seed(0   )
								snake_case_   : List[Any]						=			pipe([prompt]							,    generator=_lowercase							,    num_inference_steps=2							,    output_type="""np"""   )
								snake_case_   : Optional[int]						=			output.images

								snake_case_   : List[Any]						=			torch.Generator(device=_lowercase   ).manual_seed(0   )
								snake_case_   : Dict						=			pipe(
								    [prompt]							,    generator=_lowercase							,    output_type="""np"""							,    return_dict=_lowercase							,    num_inference_steps=2   )[0]

								snake_case_   : List[Any]						=			image[0, -3:, -3:, -1]
								snake_case_   : Any						=			image_from_tuple[0, -3:, -3:, -1]

								assert image.shape == (1, 2_4, 2_4, 3)

								snake_case_   : Dict						=			np.array([0.6551, 0.6168, 0.5008, 0.5676, 0.5659, 0.4295, 0.6073, 0.5599, 0.4992]   )

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





		def    UpperCAmelCase__							(					self   ) ->						str:







								'''simple docstring'''

								snake_case_   : int						=			"""cpu"""

								snake_case_   : List[Any]						=			self.dummy_vqvae
								snake_case_   : Optional[int]						=			self.dummy_text_encoder
								snake_case_   : List[Any]						=			self.dummy_tokenizer
								snake_case_   : Union[str, Any]						=			self.dummy_transformer
								snake_case_   : str						=			VQDiffusionScheduler(self.num_embed   )
								snake_case_   : List[Any]						=			LearnedClassifierFreeSamplingEmbeddings(
								    learnable=_lowercase							,    hidden_size=self.text_embedder_hidden_size							,    length=tokenizer.model_max_length   )

								snake_case_   : Union[str, Any]						=			VQDiffusionPipeline(
								    vqvae=_lowercase							,    text_encoder=_lowercase							,    tokenizer=_lowercase							,    transformer=_lowercase							,    scheduler=_lowercase							,    learned_classifier_free_sampling_embeddings=_lowercase							,    )
								snake_case_   : Any						=			pipe.to(_lowercase   )
								pipe.set_progress_bar_config(disable=_lowercase   )

								snake_case_   : Tuple						=			"""teddy bear playing in the pool"""

								snake_case_   : str						=			torch.Generator(device=_lowercase   ).manual_seed(0   )
								snake_case_   : Tuple						=			pipe([prompt]							,    generator=_lowercase							,    num_inference_steps=2							,    output_type="""np"""   )
								snake_case_   : Dict						=			output.images

								snake_case_   : Union[str, Any]						=			torch.Generator(device=_lowercase   ).manual_seed(0   )
								snake_case_   : Any						=			pipe(
								    [prompt]							,    generator=_lowercase							,    output_type="""np"""							,    return_dict=_lowercase							,    num_inference_steps=2   )[0]

								snake_case_   : Optional[Any]						=			image[0, -3:, -3:, -1]
								snake_case_   : int						=			image_from_tuple[0, -3:, -3:, -1]

								assert image.shape == (1, 2_4, 2_4, 3)

								snake_case_   : int						=			np.array([0.6693, 0.6075, 0.4959, 0.5701, 0.5583, 0.4333, 0.6171, 0.5684, 0.4988]   )

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




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





		"""simple docstring"""






		def    UpperCAmelCase__							(					self   ) ->						List[Any]:







								'''simple docstring'''

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





		def    UpperCAmelCase__							(					self   ) ->						List[str]:







								'''simple docstring'''

								snake_case_   : List[Any]						=			load_numpy(
								    """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"""
								    """/vq_diffusion/teddy_bear_pool_classifier_free_sampling.npy"""   )

								snake_case_   : str						=			VQDiffusionPipeline.from_pretrained("""microsoft/vq-diffusion-ithq"""   )
								snake_case_   : Optional[Any]						=			pipeline.to(_lowercase   )
								pipeline.set_progress_bar_config(disable=_lowercase   )

								# requires GPU generator for gumbel softmax
								# don't use GPU generator in tests though
								snake_case_   : Any						=			torch.Generator(device=_lowercase   ).manual_seed(0   )
								snake_case_   : Optional[int]						=			pipeline(
								    """teddy bear playing in the pool"""							,    num_images_per_prompt=1							,    generator=_lowercase							,    output_type="""np"""							,    )

								snake_case_   : Union[str, Any]						=			output.images[0]

								assert image.shape == (2_5_6, 2_5_6, 3)
								assert np.abs(expected_image - image   ).max() < 2.0






"""simple docstring"""





# Copyright (c) 2021-, NVIDIA CORPORATION.  All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

####################################################################################################

#
# Note: If when running this conversion script you're getting an exception:
#     ModuleNotFoundError: No module named 'megatron.model.enums'
# you need to tell python where to find the clone of Megatron-LM, e.g.:
#
# cd /tmp
# git clone https://github.com/NVIDIA/Megatron-LM
# PYTHONPATH=/tmp/Megatron-LM python src/transformers/models/megatron_gpt2/convert_megatron_gpt2_checkpoint.py ...
#
# if you already have it cloned elsewhere, simply adjust the path to the existing path
#
# If the training was done using a Megatron-LM fork, e.g.,
# https://github.com/microsoft/Megatron-DeepSpeed/ then chances are that you need to have that one
# in your path, i.e., /path/to/Megatron-DeepSpeed/
#

import argparse
import os
import re
import zipfile

import torch

from transformers import AutoTokenizer, GPTaConfig
def 	__lowerCAmelCase						(		__UpperCamelCase      :			Tuple					, __UpperCamelCase      :			Tuple					, __UpperCamelCase      :			Any=0    ):


						'''simple docstring'''






						if name is None:
												snake_case_   : Dict						=			None
						else:
												snake_case_   : Dict						=			""".""" * max(0					, spaces - 2    ) + """# {:""" + str(5_0 - spaces    ) + """s}"""
												snake_case_   : Any						=			fmt.format(__UpperCamelCase    )

						# Print and recurse (if needed).
						if isinstance(__UpperCamelCase					, __UpperCamelCase    ):
												if msg is not None:
																		print(__UpperCamelCase    )
												for k in val.keys():
																		recursive_print(__UpperCamelCase					, val[k]					, spaces + 2    )
						elif isinstance(__UpperCamelCase					, torch.Tensor    ):
												print(__UpperCamelCase					, """:"""					, val.size()    )
						else:
												print(__UpperCamelCase					, """:"""					, __UpperCamelCase    )
def 	__lowerCAmelCase						(		__UpperCamelCase      :			str					, __UpperCamelCase      :			List[Any]					, __UpperCamelCase      :			Dict					, __UpperCamelCase      :			str					, __UpperCamelCase      :			Any    ):


						'''simple docstring'''






						snake_case_   : Any						=			param.size()
						if checkpoint_version == 1.0:
												# version 1.0 stores [num_heads * hidden_size * num_splits, :]
												snake_case_   : List[str]						=			(num_heads, hidden_size, num_splits) + input_shape[1:]
												snake_case_   : Tuple						=			param.view(*__UpperCamelCase    )
												snake_case_   : Tuple						=			param.transpose(0					, 2    )
												snake_case_   : Any						=			param.transpose(1					, 2    ).contiguous()
						elif checkpoint_version >= 2.0:
												# other versions store [num_heads * num_splits * hidden_size, :]
												snake_case_   : Optional[Any]						=			(num_heads, num_splits, hidden_size) + input_shape[1:]
												snake_case_   : str						=			param.view(*__UpperCamelCase    )
												snake_case_   : Dict						=			param.transpose(0					, 1    ).contiguous()
						snake_case_   : int						=			param.view(*__UpperCamelCase    )
						return param
def 	__lowerCAmelCase						(		__UpperCamelCase      :			int					, __UpperCamelCase      :			Union[str, Any]					, __UpperCamelCase      :			Union[str, Any]    ):


						'''simple docstring'''






						snake_case_   : Dict						=			{}

						# old versions did not store training args
						snake_case_   : List[str]						=			input_state_dict.get("""args"""					, __UpperCamelCase    )
						if ds_args is not None:
												# do not make the user write a config file when the exact dimensions/sizes are already in the checkpoint
												# from pprint import pprint
												# pprint(vars(ds_args))

												snake_case_   : Tuple						=			ds_args.padded_vocab_size
												snake_case_   : Optional[int]						=			ds_args.max_position_embeddings
												snake_case_   : Union[str, Any]						=			ds_args.hidden_size
												snake_case_   : Union[str, Any]						=			ds_args.num_layers
												snake_case_   : str						=			ds_args.num_attention_heads
												snake_case_   : str						=			ds_args.ffn_hidden_size
												# pprint(config)

						# The number of heads.
						snake_case_   : Union[str, Any]						=			config.n_head
						# The hidden_size per head.
						snake_case_   : Optional[Any]						=			config.n_embd // config.n_head
						# Megatron-LM checkpoint version
						if "checkpoint_version" in input_state_dict.keys():
												snake_case_   : Optional[Any]						=			input_state_dict["""checkpoint_version"""]
						else:
												snake_case_   : int						=			0.0

						# The model.
						snake_case_   : List[str]						=			input_state_dict["""model"""]
						# The language model.
						snake_case_   : str						=			model["""language_model"""]
						# The embeddings.
						snake_case_   : Tuple						=			lm["""embedding"""]

						# The word embeddings.
						snake_case_   : List[str]						=			embeddings["""word_embeddings"""]["""weight"""]
						# Truncate the embedding table to vocab_size rows.
						snake_case_   : Optional[int]						=			word_embeddings[: config.vocab_size, :]
						snake_case_   : Optional[int]						=			word_embeddings

						# The position embeddings.
						snake_case_   : List[Any]						=			embeddings["""position_embeddings"""]["""weight"""]
						# Read the causal mask dimension (seqlen). [max_sequence_length, hidden_size]
						snake_case_   : Tuple						=			pos_embeddings.size(0    )
						if n_positions != config.n_positions:
												raise ValueError(
												    F'pos_embeddings.max_sequence_length={n_positions} and config.n_positions={config.n_positions} don\'t match'    )
						# Store the position embeddings.
						snake_case_   : Union[str, Any]						=			pos_embeddings

						# The transformer.
						snake_case_   : Optional[Any]						=			lm["""transformer"""] if """transformer""" in lm.keys() else lm["""encoder"""]

						# The regex to extract layer names.
						snake_case_   : Optional[Any]						=			re.compile(r"""layers\.(\d+)\.([a-z0-9_.]+)\.([a-z]+)"""    )

						# The simple map of names for "automated" rules.
						snake_case_   : List[str]						=			{
						    """attention.dense""": """.attn.c_proj.""",
						    """self_attention.dense""": """.attn.c_proj.""",
						    """mlp.dense_h_to_4h""": """.mlp.c_fc.""",
						    """mlp.dense_4h_to_h""": """.mlp.c_proj.""",
						}

						# Extract the layers.
						for key, val in transformer.items():
												# Match the name.
												snake_case_   : int						=			layer_re.match(__UpperCamelCase    )

												# Stop if that's not a layer
												if m is None:
																		break

												# The index of the layer.
												snake_case_   : Tuple						=			int(m.group(1    )    )
												# The name of the operation.
												snake_case_   : Any						=			m.group(2    )
												# Is it a weight or a bias?
												snake_case_   : Union[str, Any]						=			m.group(3    )

												# The name of the layer.
												snake_case_   : str						=			F'transformer.h.{layer_idx}'

												# For layernorm(s), simply store the layer norm.
												if op_name.endswith("""layernorm"""    ):
																		snake_case_   : Dict						=			"""ln_1""" if op_name.startswith("""input"""    ) else """ln_2"""
																		snake_case_   : Optional[int]						=			val

												# Transpose the QKV matrix.
												elif (
												    op_name == "attention.query_key_value" or op_name == "self_attention.query_key_value"
												) and weight_or_bias == "weight":
																		# Insert a tensor of 1x1xDxD bias.
																		snake_case_   : Optional[Any]						=			torch.tril(torch.ones((n_positions, n_positions)					, dtype=torch.floataa    )    ).view(
																		    1					, 1					, __UpperCamelCase					, __UpperCamelCase    )
																		snake_case_   : List[Any]						=			causal_mask

																		# Insert a "dummy" tensor for masked_bias.
																		snake_case_   : str						=			torch.tensor(-1E4					, dtype=torch.floataa    )
																		snake_case_   : List[Any]						=			masked_bias

																		snake_case_   : Optional[int]						=			fix_query_key_value_ordering(__UpperCamelCase					, __UpperCamelCase					, 3					, __UpperCamelCase					, __UpperCamelCase    )
																		# Megatron stores (3*D) x D but transformers-GPT2 expects D x 3*D.
																		snake_case_   : str						=			out_val.transpose(0					, 1    ).contiguous()
																		# Store.
																		snake_case_   : Tuple						=			out_val

												# Transpose the bias.
												elif (
												    op_name == "attention.query_key_value" or op_name == "self_attention.query_key_value"
												) and weight_or_bias == "bias":
																		snake_case_   : Optional[Any]						=			fix_query_key_value_ordering(__UpperCamelCase					, __UpperCamelCase					, 3					, __UpperCamelCase					, __UpperCamelCase    )
																		# Store. No change of shape.
																		snake_case_   : List[Any]						=			out_val

												# Transpose the weights.
												elif weight_or_bias == "weight":
																		snake_case_   : Any						=			megatron_to_transformers[op_name]
																		snake_case_   : str						=			val.transpose(0					, 1    )

												# Copy the bias.
												elif weight_or_bias == "bias":
																		snake_case_   : List[str]						=			megatron_to_transformers[op_name]
																		snake_case_   : Tuple						=			val

    # DEBUG.
						assert config.n_layer == layer_idx + 1

						# The final layernorm.
						snake_case_   : Dict						=			transformer["""final_layernorm.weight"""]
						snake_case_   : Dict						=			transformer["""final_layernorm.bias"""]

						# For LM head, transformers' wants the matrix to weight embeddings.
						snake_case_   : Optional[int]						=			word_embeddings

						# It should be done!
						return output_state_dict
def 	__lowerCAmelCase						(		):


						'''simple docstring'''






						snake_case_   : List[str]						=			argparse.ArgumentParser()
						parser.add_argument("""--print-checkpoint-structure"""					, action="""store_true"""    )
						parser.add_argument(
						    """path_to_checkpoint"""					, type=__UpperCamelCase					, help="""Path to the checkpoint file (.zip archive or direct .pt file)"""					, )
						parser.add_argument(
						    """--config_file"""					, default=""""""					, type=__UpperCamelCase					, help="""An optional config json file describing the pre-trained model."""					, )
						snake_case_   : str						=			parser.parse_args()

						# Extract the basename.
						snake_case_   : Optional[Any]						=			os.path.dirname(args.path_to_checkpoint    )

						# Load the model.
						# the .zip is very optional, let's keep it for backward compatibility
						print(F'Extracting PyTorch state dictionary from {args.path_to_checkpoint}'    )
						if args.path_to_checkpoint.endswith(""".zip"""    ):
												with zipfile.ZipFile(args.path_to_checkpoint					, """r"""    ) as checkpoint:
																		with checkpoint.open("""release/mp_rank_00/model_optim_rng.pt"""    ) as pytorch_dict:
																								snake_case_   : Optional[int]						=			torch.load(__UpperCamelCase					, map_location="""cpu"""    )
						else:
												snake_case_   : List[Any]						=			torch.load(args.path_to_checkpoint					, map_location="""cpu"""    )

						snake_case_   : Any						=			input_state_dict.get("""args"""					, __UpperCamelCase    )

						# Read the config, or default to the model released by NVIDIA.
						if args.config_file == "":
												if ds_args is not None:
																		if ds_args.bias_gelu_fusion:
																								snake_case_   : Any						=			"""gelu_fast"""
																		elif ds_args.openai_gelu:
																								snake_case_   : Tuple						=			"""gelu_new"""
																		else:
																								snake_case_   : List[str]						=			"""gelu"""
												else:
																		# in the very early days this used to be "gelu_new"
																		snake_case_   : Dict						=			"""gelu_new"""

												# Spell out all parameters in case the defaults change.
												snake_case_   : List[str]						=			GPTaConfig(
												    vocab_size=5_0_2_5_7					, n_positions=1_0_2_4					, n_embd=1_0_2_4					, n_layer=2_4					, n_head=1_6					, n_inner=4_0_9_6					, activation_function=__UpperCamelCase					, resid_pdrop=0.1					, embd_pdrop=0.1					, attn_pdrop=0.1					, layer_norm_epsilon=1E-5					, initializer_range=0.02					, summary_type="""cls_index"""					, summary_use_proj=__UpperCamelCase					, summary_activation=__UpperCamelCase					, summary_proj_to_labels=__UpperCamelCase					, summary_first_dropout=0.1					, scale_attn_weights=__UpperCamelCase					, use_cache=__UpperCamelCase					, bos_token_id=5_0_2_5_6					, eos_token_id=5_0_2_5_6					, )
						else:
												snake_case_   : List[Any]						=			GPTaConfig.from_json_file(args.config_file    )

						snake_case_   : int						=			["""GPT2LMHeadModel"""]

						# Convert.
						print("""Converting"""    )
						snake_case_   : Tuple						=			convert_megatron_checkpoint(__UpperCamelCase					, __UpperCamelCase					, __UpperCamelCase    )

						# Print the structure of converted state dict.
						if args.print_checkpoint_structure:
												recursive_print(__UpperCamelCase					, __UpperCamelCase    )

						# Add tokenizer class info to config
						# see https://github.com/huggingface/transformers/issues/13906)
						if ds_args is not None:
												snake_case_   : str						=			ds_args.tokenizer_type
												if tokenizer_type == "GPT2BPETokenizer":
																		snake_case_   : Optional[Any]						=			"""gpt2"""
												elif tokenizer_type == "PretrainedFromHF":
																		snake_case_   : str						=			ds_args.tokenizer_name_or_path
												else:
																		raise ValueError(F'Unrecognized tokenizer_type {tokenizer_type}'    )
						else:
												snake_case_   : List[str]						=			"""gpt2"""

						snake_case_   : List[Any]						=			AutoTokenizer.from_pretrained(__UpperCamelCase    )
						snake_case_   : List[str]						=			type(__UpperCamelCase    ).__name__
						snake_case_   : Optional[int]						=			tokenizer_class

						# Store the config to file.
						print("""Saving config"""    )
						config.save_pretrained(__UpperCamelCase    )

						# Save tokenizer based on args
						print(F'Adding {tokenizer_class} tokenizer files'    )
						tokenizer.save_pretrained(__UpperCamelCase    )

						# Store the state_dict to file.
						snake_case_   : List[Any]						=			os.path.join(__UpperCamelCase					, """pytorch_model.bin"""    )
						print(F'Saving checkpoint to "{output_checkpoint_file}"'    )
						torch.save(__UpperCamelCase					, __UpperCamelCase    )


####################################################################################################

if __name__ == "__main__":
					main()

####################################################################################################

