


import argparse
import os
import re
import zipfile

import torch

from transformers import AutoTokenizer, GPTaConfig




def      UpperCAmelCase__ (lowerCAmelCase_						,		lowerCAmelCase_						,		lowerCAmelCase_=0     ):




					'''simple docstring'''


					if name is None:
										__SCREAMING_SNAKE_CASE		     =       None
					else:
										__SCREAMING_SNAKE_CASE		     =       "." * max(0						,		spaces - 2     ) + "# {:" + str(50 - spaces     ) + "s}"
										__SCREAMING_SNAKE_CASE		     =       fmt.format(lowerCAmelCase_     )

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




def      UpperCAmelCase__ (lowerCAmelCase_						,		lowerCAmelCase_						,		lowerCAmelCase_						,		lowerCAmelCase_						,		lowerCAmelCase_     ):




					'''simple docstring'''


					__SCREAMING_SNAKE_CASE		     =       param.size()
					if checkpoint_version == 1.0:
										# version 1.0 stores [num_heads * hidden_size * num_splits, :]
										__SCREAMING_SNAKE_CASE		     =       (num_heads, hidden_size, num_splits) + input_shape[1:]
										__SCREAMING_SNAKE_CASE		     =       param.view(*lowerCAmelCase_     )
										__SCREAMING_SNAKE_CASE		     =       param.transpose(0						,		2     )
										__SCREAMING_SNAKE_CASE		     =       param.transpose(1						,		2     ).contiguous()
					elif checkpoint_version >= 2.0:
										# other versions store [num_heads * num_splits * hidden_size, :]
										__SCREAMING_SNAKE_CASE		     =       (num_heads, num_splits, hidden_size) + input_shape[1:]
										__SCREAMING_SNAKE_CASE		     =       param.view(*lowerCAmelCase_     )
										__SCREAMING_SNAKE_CASE		     =       param.transpose(0						,		1     ).contiguous()
					__SCREAMING_SNAKE_CASE		     =       param.view(*lowerCAmelCase_     )
					return param




def      UpperCAmelCase__ (lowerCAmelCase_						,		lowerCAmelCase_						,		lowerCAmelCase_     ):




					'''simple docstring'''


					__SCREAMING_SNAKE_CASE		     =       {}

					# old versions did not store training args
					__SCREAMING_SNAKE_CASE		     =       input_state_dict.get("args"						,		lowerCAmelCase_     )
					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))

										__SCREAMING_SNAKE_CASE		     =       ds_args.padded_vocab_size
										__SCREAMING_SNAKE_CASE		     =       ds_args.max_position_embeddings
										__SCREAMING_SNAKE_CASE		     =       ds_args.hidden_size
										__SCREAMING_SNAKE_CASE		     =       ds_args.num_layers
										__SCREAMING_SNAKE_CASE		     =       ds_args.num_attention_heads
										__SCREAMING_SNAKE_CASE		     =       ds_args.ffn_hidden_size
										# pprint(config)

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

					# The model.
					__SCREAMING_SNAKE_CASE		     =       input_state_dict["model"]
					# The language model.
					__SCREAMING_SNAKE_CASE		     =       model["language_model"]
					# The embeddings.
					__SCREAMING_SNAKE_CASE		     =       lm["embedding"]

					# The word embeddings.
					__SCREAMING_SNAKE_CASE		     =       embeddings["word_embeddings"]["weight"]
					# Truncate the embedding table to vocab_size rows.
					__SCREAMING_SNAKE_CASE		     =       word_embeddings[: config.vocab_size, :]
					__SCREAMING_SNAKE_CASE		     =       word_embeddings

					# The position embeddings.
					__SCREAMING_SNAKE_CASE		     =       embeddings["position_embeddings"]["weight"]
					# Read the causal mask dimension (seqlen). [max_sequence_length, hidden_size]
					__SCREAMING_SNAKE_CASE		     =       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.
					__SCREAMING_SNAKE_CASE		     =       pos_embeddings

					# The transformer.
					__SCREAMING_SNAKE_CASE		     =       lm["transformer"] if "transformer" in lm.keys() else lm["encoder"]

					# The regex to extract layer names.
					__SCREAMING_SNAKE_CASE		     =       re.compile(R"layers\.(\d+)\.([a-z0-9_.]+)\.([a-z]+)"     )

					# The simple map of names for "automated" rules.
					__SCREAMING_SNAKE_CASE		     =       {
					    "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.
										__SCREAMING_SNAKE_CASE		     =       layer_re.match(lowerCAmelCase_     )

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

										# The index of the layer.
										__SCREAMING_SNAKE_CASE		     =       int(m.group(1     )     )
										# The name of the operation.
										__SCREAMING_SNAKE_CASE		     =       m.group(2     )
										# Is it a weight or a bias?
										__SCREAMING_SNAKE_CASE		     =       m.group(3     )

										# The name of the layer.
										__SCREAMING_SNAKE_CASE		     =       f"""transformer.h.{layer_idx}"""

										# For layernorm(s), simply store the layer norm.
										if op_name.endswith("layernorm"     ):
															__SCREAMING_SNAKE_CASE		     =       "ln_1" if op_name.startswith("input"     ) else "ln_2"
															__SCREAMING_SNAKE_CASE		     =       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.
															__SCREAMING_SNAKE_CASE		     =       torch.tril(torch.ones((n_positions, n_positions)						,		dtype=torch.floataa     )     ).view(
															    1						,		1						,		lowerCAmelCase_						,		lowerCAmelCase_     )
															__SCREAMING_SNAKE_CASE		     =       causal_mask

															# Insert a "dummy" tensor for masked_bias.
															__SCREAMING_SNAKE_CASE		     =       torch.tensor(-1E4						,		dtype=torch.floataa     )
															__SCREAMING_SNAKE_CASE		     =       masked_bias

															__SCREAMING_SNAKE_CASE		     =       fix_query_key_value_ordering(lowerCAmelCase_						,		lowerCAmelCase_						,		3						,		lowerCAmelCase_						,		lowerCAmelCase_     )
															# Megatron stores (3*D) x D but transformers-GPT2 expects D x 3*D.
															__SCREAMING_SNAKE_CASE		     =       out_val.transpose(0						,		1     ).contiguous()
															# Store.
															__SCREAMING_SNAKE_CASE		     =       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":
															__SCREAMING_SNAKE_CASE		     =       fix_query_key_value_ordering(lowerCAmelCase_						,		lowerCAmelCase_						,		3						,		lowerCAmelCase_						,		lowerCAmelCase_     )
															# Store. No change of shape.
															__SCREAMING_SNAKE_CASE		     =       out_val

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

										# Copy the bias.
										elif weight_or_bias == "bias":
															__SCREAMING_SNAKE_CASE		     =       megatron_to_transformers[op_name]
															__SCREAMING_SNAKE_CASE		     =       val

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

					# The final layernorm.
					__SCREAMING_SNAKE_CASE		     =       transformer["final_layernorm.weight"]
					__SCREAMING_SNAKE_CASE		     =       transformer["final_layernorm.bias"]

					# For LM head, transformers' wants the matrix to weight embeddings.
					__SCREAMING_SNAKE_CASE		     =       word_embeddings

					# It should be done!
					return output_state_dict




def      UpperCAmelCase__ ():




					'''simple docstring'''


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

					# Extract the basename.
					__SCREAMING_SNAKE_CASE		     =       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:
																				__SCREAMING_SNAKE_CASE		     =       torch.load(lowerCAmelCase_						,		map_location="cpu"     )
					else:
										__SCREAMING_SNAKE_CASE		     =       torch.load(args.path_to_checkpoint						,		map_location="cpu"     )

					__SCREAMING_SNAKE_CASE		     =       input_state_dict.get("args"						,		lowerCAmelCase_     )

					# 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:
																				__SCREAMING_SNAKE_CASE		     =       "gelu_fast"
															elif ds_args.openai_gelu:
																				__SCREAMING_SNAKE_CASE		     =       "gelu_new"
															else:
																				__SCREAMING_SNAKE_CASE		     =       "gelu"
										else:
															# in the very early days this used to be "gelu_new"
															__SCREAMING_SNAKE_CASE		     =       "gelu_new"

										# Spell out all parameters in case the defaults change.
										__SCREAMING_SNAKE_CASE		     =       GPTaConfig(
										    vocab_size=5_0257						,		n_positions=1024						,		n_embd=1024						,		n_layer=24						,		n_head=16						,		n_inner=4096						,		activation_function=lowerCAmelCase_						,		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=lowerCAmelCase_						,		summary_activation=lowerCAmelCase_						,		summary_proj_to_labels=lowerCAmelCase_						,		summary_first_dropout=0.1						,		scale_attn_weights=lowerCAmelCase_						,		use_cache=lowerCAmelCase_						,		bos_token_id=5_0256						,		eos_token_id=5_0256						,		)
					else:
										__SCREAMING_SNAKE_CASE		     =       GPTaConfig.from_json_file(args.config_file     )

					__SCREAMING_SNAKE_CASE		     =       ["GPT2LMHeadModel"]

					# Convert.
					print("Converting"     )
					__SCREAMING_SNAKE_CASE		     =       convert_megatron_checkpoint(lowerCAmelCase_						,		lowerCAmelCase_						,		lowerCAmelCase_     )

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

					# Add tokenizer class info to config
					# see https://github.com/huggingface/transformers/issues/13906)
					if ds_args is not None:
										__SCREAMING_SNAKE_CASE		     =       ds_args.tokenizer_type
										if tokenizer_type == "GPT2BPETokenizer":
															__SCREAMING_SNAKE_CASE		     =       "gpt2"
										elif tokenizer_type == "PretrainedFromHF":
															__SCREAMING_SNAKE_CASE		     =       ds_args.tokenizer_name_or_path
										else:
															raise ValueError(f"""Unrecognized tokenizer_type {tokenizer_type}"""     )
					else:
										__SCREAMING_SNAKE_CASE		     =       "gpt2"

					__SCREAMING_SNAKE_CASE		     =       AutoTokenizer.from_pretrained(lowerCAmelCase_     )
					__SCREAMING_SNAKE_CASE		     =       type(lowerCAmelCase_     ).__name__
					__SCREAMING_SNAKE_CASE		     =       tokenizer_class

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

					# Save tokenizer based on args
					print(f"""Adding {tokenizer_class} tokenizer files"""     )
					tokenizer.save_pretrained(lowerCAmelCase_     )

					# Store the state_dict to file.
					__SCREAMING_SNAKE_CASE		     =       os.path.join(lowerCAmelCase_						,		"pytorch_model.bin"     )
					print(f"""Saving checkpoint to \"{output_checkpoint_file}\""""     )
					torch.save(lowerCAmelCase_						,		lowerCAmelCase_     )


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

if __name__ == "__main__":
		main()

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




"""simple docstring"""







import gc
import unittest

import numpy as np
import torch

from diffusers import AutoencoderKL, DDIMScheduler, DiTPipeline, DPMSolverMultistepScheduler, TransformeraDModel
from diffusers.utils import is_xformers_available, load_numpy, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..pipeline_params import (
    CLASS_CONDITIONED_IMAGE_GENERATION_BATCH_PARAMS,
    CLASS_CONDITIONED_IMAGE_GENERATION_PARAMS,
)
from ..test_pipelines_common import PipelineTesterMixin


enable_full_determinism()




class 				UpperCamelCase_			(				UpperCamelCase    ,						unittest.TestCase):

	"""simple docstring"""





	snake_case__       :						Union[str, Any]					      =    DiTPipeline
	snake_case__       :						Union[str, Any]					      =    CLASS_CONDITIONED_IMAGE_GENERATION_PARAMS
	snake_case__       :						Any					      =    PipelineTesterMixin.required_optional_params - {
	    "latents",
	    "num_images_per_prompt",
	    "callback",
	    "callback_steps",
	}
	snake_case__       :						Optional[Any]					      =    CLASS_CONDITIONED_IMAGE_GENERATION_BATCH_PARAMS
	snake_case__       :						Optional[Any]					      =    False




	def 						UpperCAmelCase_   (				self    :       Dict     ) ->			Dict:
						torch.manual_seed(0     )
						__SCREAMING_SNAKE_CASE		     =       TransformeraDModel(
						    sample_size=1_6    ,    num_layers=2    ,    patch_size=4    ,    attention_head_dim=8    ,    num_attention_heads=2    ,    in_channels=4    ,    out_channels=8    ,    attention_bias=UpperCAmelCase__    ,    activation_fn="gelu-approximate"    ,    num_embeds_ada_norm=1_0_0_0    ,    norm_type="ada_norm_zero"    ,    norm_elementwise_affine=UpperCAmelCase__    ,    )
						__SCREAMING_SNAKE_CASE		     =       AutoencoderKL()
						__SCREAMING_SNAKE_CASE		     =       DDIMScheduler()
						__SCREAMING_SNAKE_CASE		     =       {"transformer": transformer.eval(), "vae": vae.eval(), "scheduler": scheduler}
						return components




	def 						UpperCAmelCase_   (				self    :       Optional[int]    ,    UpperCAmelCase__    :       Tuple    ,    UpperCAmelCase__    :       str=0     ) ->			Optional[int]:
						if str(UpperCAmelCase__     ).startswith("mps"     ):
											__SCREAMING_SNAKE_CASE		     =       torch.manual_seed(UpperCAmelCase__     )
						else:
											__SCREAMING_SNAKE_CASE		     =       torch.Generator(device=UpperCAmelCase__     ).manual_seed(UpperCAmelCase__     )
						__SCREAMING_SNAKE_CASE		     =       {
						    "class_labels": [1],
						    "generator": generator,
						    "num_inference_steps": 2,
						    "output_type": "numpy",
						}
						return inputs




	def 						UpperCAmelCase_   (				self    :       Optional[int]     ) ->			List[str]:
						__SCREAMING_SNAKE_CASE		     =       "cpu"

						__SCREAMING_SNAKE_CASE		     =       self.get_dummy_components()
						__SCREAMING_SNAKE_CASE		     =       self.pipeline_class(**UpperCAmelCase__     )
						pipe.to(UpperCAmelCase__     )
						pipe.set_progress_bar_config(disable=UpperCAmelCase__     )

						__SCREAMING_SNAKE_CASE		     =       self.get_dummy_inputs(UpperCAmelCase__     )
						__SCREAMING_SNAKE_CASE		     =       pipe(**UpperCAmelCase__     ).images
						__SCREAMING_SNAKE_CASE		     =       image[0, -3:, -3:, -1]

						self.assertEqual(image.shape    ,    (1, 1_6, 1_6, 3)     )
						__SCREAMING_SNAKE_CASE		     =       np.array([0.2_946, 0.6_601, 0.4_329, 0.3_296, 0.4_144, 0.5_319, 0.7_273, 0.5_013, 0.4_457]     )
						__SCREAMING_SNAKE_CASE		     =       np.abs(image_slice.flatten() - expected_slice     ).max()
						self.assertLessEqual(UpperCAmelCase__    ,    1E-3     )




	def 						UpperCAmelCase_   (				self    :       List[str]     ) ->			Union[str, Any]:
						self._test_inference_batch_single_identical(relax_max_difference=UpperCAmelCase__    ,    expected_max_diff=1E-3     )




	@unittest.skipIf(
	    torch_device != "cuda" or not is_xformers_available()    ,    reason="XFormers attention is only available with CUDA and `xformers` installed"    ,    )
	def 						UpperCAmelCase_   (				self    :       int     ) ->			Dict:
						self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=1E-3     )




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

	"""simple docstring"""





	def 						UpperCAmelCase_   (				self    :       Optional[int]     ) ->			Any:
						super().tearDown()
						gc.collect()
						torch.cuda.empty_cache()




	def 						UpperCAmelCase_   (				self    :       Union[str, Any]     ) ->			Optional[int]:
						__SCREAMING_SNAKE_CASE		     =       torch.manual_seed(0     )

						__SCREAMING_SNAKE_CASE		     =       DiTPipeline.from_pretrained("facebook/DiT-XL-2-256"     )
						pipe.to("cuda"     )

						__SCREAMING_SNAKE_CASE		     =       ["vase", "umbrella", "white shark", "white wolf"]
						__SCREAMING_SNAKE_CASE		     =       pipe.get_label_ids(UpperCAmelCase__     )

						__SCREAMING_SNAKE_CASE		     =       pipe(UpperCAmelCase__    ,    generator=UpperCAmelCase__    ,    num_inference_steps=4_0    ,    output_type="np"     ).images

						for word, image in zip(UpperCAmelCase__    ,    UpperCAmelCase__     ):
											__SCREAMING_SNAKE_CASE		     =       load_numpy(
											    F"""https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/dit/{word}.npy"""     )
											assert np.abs((expected_image - image).max()     ) < 1E-2




	def 						UpperCAmelCase_   (				self    :       List[str]     ) ->			List[Any]:
						__SCREAMING_SNAKE_CASE		     =       DiTPipeline.from_pretrained("facebook/DiT-XL-2-512"     )
						__SCREAMING_SNAKE_CASE		     =       DPMSolverMultistepScheduler.from_config(pipe.scheduler.config     )
						pipe.to("cuda"     )

						__SCREAMING_SNAKE_CASE		     =       ["vase", "umbrella"]
						__SCREAMING_SNAKE_CASE		     =       pipe.get_label_ids(UpperCAmelCase__     )

						__SCREAMING_SNAKE_CASE		     =       torch.manual_seed(0     )
						__SCREAMING_SNAKE_CASE		     =       pipe(UpperCAmelCase__    ,    generator=UpperCAmelCase__    ,    num_inference_steps=2_5    ,    output_type="np"     ).images

						for word, image in zip(UpperCAmelCase__    ,    UpperCAmelCase__     ):
											__SCREAMING_SNAKE_CASE		     =       load_numpy(
											    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
											    F"""/dit/{word}_512.npy"""     )

											assert np.abs((expected_image - image).max()     ) < 1E-1
