

'''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 lowerCamelCase			(  lowercase_	,						lowercase_	,						lowercase_	,						unittest.TestCase						):



			'''simple docstring'''






			__snake_case        =	AltDiffusionPipeline
			__snake_case        =	TEXT_TO_IMAGE_PARAMS
			__snake_case        =	TEXT_TO_IMAGE_BATCH_PARAMS
			__snake_case        =	TEXT_TO_IMAGE_IMAGE_PARAMS
			__snake_case        =	TEXT_TO_IMAGE_IMAGE_PARAMS


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





									'''simple docstring'''
									torch.manual_seed(0			)
									A__					:       List[Any]     =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		,      )
									A__					:       Union[str, Any]     =DDIMScheduler(
									    beta_start=0.00085		,      beta_end=0.012		,      beta_schedule="""scaled_linear"""		,      clip_sample=lowerCAmelCase_		,      set_alpha_to_one=lowerCAmelCase_		,      )
									torch.manual_seed(0			)
									A__					:       Dict     =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			)
									A__					:       Optional[int]     =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		,      )
									A__					:       Tuple     =CLIPTextModel(lowerCAmelCase_			)

									A__					:       Dict     =XLMRobertaTokenizer.from_pretrained("""hf-internal-testing/tiny-xlm-roberta"""			)
									A__					:       List[str]     =77

									A__					:       Tuple     ={
									    """unet""": unet,
									    """scheduler""": scheduler,
									    """vae""": vae,
									    """text_encoder""": text_encoder,
									    """tokenizer""": tokenizer,
									    """safety_checker""": None,
									    """feature_extractor""": None,
									}
									return components


			def lowercase__						(				self	:						List[Any]		,      lowerCAmelCase_	:						Tuple		,      lowerCAmelCase_	:						Optional[int]=0			)      ->		Optional[int]:





									'''simple docstring'''
									if str(lowerCAmelCase_			).startswith("""mps"""			):
															A__					:       Optional[Any]     =torch.manual_seed(lowerCAmelCase_			)
									else:
															A__					:       Any     =torch.Generator(device=lowerCAmelCase_			).manual_seed(lowerCAmelCase_			)
									A__					:       int     ={
									    """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 lowercase__						(				self	:						str			)      ->		str:





									'''simple docstring'''
									super().test_attention_slicing_forward_pass(expected_max_diff=3e-3			)


			def lowercase__						(				self	:						List[Any]			)      ->		Optional[Any]:





									'''simple docstring'''
									super().test_inference_batch_single_identical(expected_max_diff=3e-3			)


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





									'''simple docstring'''
									A__					:       str     ="""cpu"""  # ensure determinism for the device-dependent torch.Generator

									A__					:       Union[str, Any]     =self.get_dummy_components()
									torch.manual_seed(0			)
									A__					:       Any     =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
									A__					:       int     =RobertaSeriesModelWithTransformation(lowerCAmelCase_			)
									A__					:       int     =text_encoder

									A__					:       Dict     =AltDiffusionPipeline(**lowerCAmelCase_			)
									A__					:       Dict     =alt_pipe.to(lowerCAmelCase_			)
									alt_pipe.set_progress_bar_config(disable=lowerCAmelCase_			)

									A__					:       List[str]     =self.get_dummy_inputs(lowerCAmelCase_			)
									A__					:       Dict     ="""A photo of an astronaut"""
									A__					:       Dict     =alt_pipe(**lowerCAmelCase_			)
									A__					:       List[Any]     =output.images
									A__					:       List[Any]     =image[0, -3:, -3:, -1]

									assert image.shape == (1, 64, 64, 3)
									A__					:       Union[str, Any]     =np.array(
									    [0.5748162, 0.60447145, 0.48821217, 0.50100636, 0.5431185, 0.45763683, 0.49657696, 0.48132733, 0.47573093]			)

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



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





									'''simple docstring'''
									A__					:       int     ="""cpu"""  # ensure determinism for the device-dependent torch.Generator

									A__					:       Any     =self.get_dummy_components()
									A__					:       str     =PNDMScheduler(skip_prk_steps=lowerCAmelCase_			)
									torch.manual_seed(0			)
									A__					:       Any     =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
									A__					:       List[Any]     =RobertaSeriesModelWithTransformation(lowerCAmelCase_			)
									A__					:       List[Any]     =text_encoder
									A__					:       int     =AltDiffusionPipeline(**lowerCAmelCase_			)
									A__					:       Union[str, Any]     =alt_pipe.to(lowerCAmelCase_			)
									alt_pipe.set_progress_bar_config(disable=lowerCAmelCase_			)

									A__					:       Dict     =self.get_dummy_inputs(lowerCAmelCase_			)
									A__					:       List[Any]     =alt_pipe(**lowerCAmelCase_			)
									A__					:       Dict     =output.images
									A__					:       str     =image[0, -3:, -3:, -1]

									assert image.shape == (1, 64, 64, 3)
									A__					:       Optional[int]     =np.array(
									    [0.51605093, 0.5707241, 0.47365507, 0.50578886, 0.5633877, 0.4642503, 0.5182081, 0.48763484, 0.49084237]			)

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



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



			'''simple docstring'''






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





									'''simple docstring'''
									# clean up the VRAM after each test
									super().tearDown()
									gc.collect()
									torch.cuda.empty_cache()


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





									'''simple docstring'''
									# make sure here that pndm scheduler skips prk
									A__					:       List[str]     =AltDiffusionPipeline.from_pretrained("""BAAI/AltDiffusion"""		,      safety_checker=lowerCAmelCase_			)
									A__					:       List[str]     =alt_pipe.to(lowerCAmelCase_			)
									alt_pipe.set_progress_bar_config(disable=lowerCAmelCase_			)

									A__					:       Union[str, Any]     ="""A painting of a squirrel eating a burger"""
									A__					:       Optional[int]     =torch.manual_seed(0			)
									A__					:       Optional[Any]     =alt_pipe([prompt]		,      generator=lowerCAmelCase_		,      guidance_scale=6.0		,      num_inference_steps=20		,      output_type="""np"""			)

									A__					:       Tuple     =output.images

									A__					:       Union[str, Any]     =image[0, -3:, -3:, -1]

									assert image.shape == (1, 5_12, 5_12, 3)
									A__					:       Tuple     =np.array([0.1010, 0.0800, 0.0794, 0.0885, 0.0843, 0.0762, 0.0769, 0.0729, 0.0586]			)

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



			def lowercase__						(				self	:						Tuple			)      ->		Dict:





									'''simple docstring'''
									A__					:       Any     =DDIMScheduler.from_pretrained("""BAAI/AltDiffusion"""		,      subfolder="""scheduler"""			)

									A__					:       List[str]     =AltDiffusionPipeline.from_pretrained("""BAAI/AltDiffusion"""		,      scheduler=lowerCAmelCase_		,      safety_checker=lowerCAmelCase_			)
									A__					:       int     =alt_pipe.to(lowerCAmelCase_			)
									alt_pipe.set_progress_bar_config(disable=lowerCAmelCase_			)

									A__					:       Optional[int]     ="""A painting of a squirrel eating a burger"""
									A__					:       List[Any]     =torch.manual_seed(0			)

									A__					:       List[Any]     =alt_pipe([prompt]		,      generator=lowerCAmelCase_		,      num_inference_steps=2		,      output_type="""numpy"""			)
									A__					:       List[Any]     =output.images

									A__					:       Optional[int]     =image[0, -3:, -3:, -1]

									assert image.shape == (1, 5_12, 5_12, 3)
									A__					:       List[str]     =np.array([0.4019, 0.4052, 0.3810, 0.4119, 0.3916, 0.3982, 0.4651, 0.4195, 0.5323]			)

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




'''simple docstring'''



import argparse
import gc
import json
import os
import shutil
import warnings

import torch

from transformers import LlamaConfig, LlamaForCausalLM, LlamaTokenizer


try:
		from transformers import LlamaTokenizerFast
except ImportError as e:
		warnings.warn(e)
		warnings.warn(
		    'The converted tokenizer will be the `slow` tokenizer. To use the fast, update your `tokenizers` library and re-run the tokenizer conversion'
		)
		__snake_case			:     Optional[Any]										=      None
__snake_case			:     Optional[Any]										=      {
    '7B': 1_1008,
    '13B': 1_3824,
    '30B': 1_7920,
    '65B': 2_2016,
    '70B': 2_8672,
}
__snake_case			:     Union[str, Any]										=      {
    '7B': 1,
    '7Bf': 1,
    '13B': 2,
    '13Bf': 2,
    '30B': 4,
    '65B': 8,
    '70B': 8,
    '70Bf': 8,
}
def 						__lowerCamelCase	(       __snake_case					:		Optional[Any],		__snake_case					:		str=1,		__snake_case					:		Tuple=256	)	->							str:






						"""simple docstring"""





						return multiple_of * ((int(ffn_dim_multiplier * int(8 * n / 3	)	) + multiple_of - 1) // multiple_of)
def 						__lowerCamelCase	(       __snake_case					:		Tuple	)	->							Tuple:






						"""simple docstring"""





						with open(__snake_case,		"""r"""	) as f:
												return json.load(__snake_case	)
def 						__lowerCamelCase	(       __snake_case					:		Optional[int],		__snake_case					:		Tuple	)	->							Dict:






						"""simple docstring"""





						with open(__snake_case,		"""w"""	) as f:
												json.dump(__snake_case,		__snake_case	)
def 						__lowerCamelCase	(       __snake_case					:		List[Any],		__snake_case					:		Any,		__snake_case					:		Any,		__snake_case					:		Tuple=True	)	->							Any:






						"""simple docstring"""





						os.makedirs(__snake_case,		exist_ok=__snake_case	)
						A__					:       List[Any]     =os.path.join(__snake_case,		"""tmp"""	)
						os.makedirs(__snake_case,		exist_ok=__snake_case	)

						A__					:       Dict     =read_json(os.path.join(__snake_case,		"""params.json"""	)	)
						A__					:       Dict     =NUM_SHARDS[model_size]
						A__					:       List[str]     =params["""n_layers"""]
						A__					:       int     =params["""n_heads"""]
						A__					:       str     =n_heads // num_shards
						A__					:       Tuple     =params["""dim"""]
						A__					:       Union[str, Any]     =dim // n_heads
						A__					:       str     =1_00_00.0
						A__					:       Any     =1.0 / (base ** (torch.arange(0,		__snake_case,		2	).float() / dims_per_head))

						if "n_kv_heads" in params:
												A__					:       Optional[Any]     =params["""n_kv_heads"""]  # for GQA / MQA
												A__					:       int     =n_heads_per_shard // num_key_value_heads
												A__					:       int     =dim // num_key_value_heads
						else:  # compatibility with other checkpoints
												A__					:       List[Any]     =n_heads
												A__					:       List[str]     =n_heads_per_shard
												A__					:       Dict     =dim

						# permute for sliced rotary
						def permute(__snake_case					:		Tuple,		__snake_case					:		Optional[int]=n_heads,		__snake_case					:		int=dim,		__snake_case					:		Optional[Any]=dim	):
												return w.view(__snake_case,		dima // n_heads // 2,		2,		__snake_case	).transpose(1,		2	).reshape(__snake_case,		__snake_case	)

						print(f"Fetching all parameters from the checkpoint at {input_base_path}."	)
						# Load weights
						if model_size == "7B":
												# Not sharded
												# (The sharded implementation would also work, but this is simpler.)
												A__					:       List[str]     =torch.load(os.path.join(__snake_case,		"""consolidated.00.pth"""	),		map_location="""cpu"""	)
						else:
												# Sharded
												A__					:       Optional[Any]     =[
												    torch.load(os.path.join(__snake_case,		f"consolidated.{i:02d}.pth"	),		map_location="""cpu"""	)
												    for i in range(__snake_case	)
												]
						A__					:       Optional[Any]     =0
						A__					:       str     ={"""weight_map""": {}}
						for layer_i in range(__snake_case	):
												A__					:       Dict     =f"pytorch_model-{layer_i + 1}-of-{n_layers + 1}.bin"
												if model_size == "7B":
																		# Unsharded
																		A__					:       Dict     ={
																		    f"model.layers.{layer_i}.self_attn.q_proj.weight": permute(
																		        loaded[f"layers.{layer_i}.attention.wq.weight"]	),
																		    f"model.layers.{layer_i}.self_attn.k_proj.weight": permute(
																		        loaded[f"layers.{layer_i}.attention.wk.weight"]	),
																		    f"model.layers.{layer_i}.self_attn.v_proj.weight": loaded[f"layers.{layer_i}.attention.wv.weight"],
																		    f"model.layers.{layer_i}.self_attn.o_proj.weight": loaded[f"layers.{layer_i}.attention.wo.weight"],
																		    f"model.layers.{layer_i}.mlp.gate_proj.weight": loaded[f"layers.{layer_i}.feed_forward.w1.weight"],
																		    f"model.layers.{layer_i}.mlp.down_proj.weight": loaded[f"layers.{layer_i}.feed_forward.w2.weight"],
																		    f"model.layers.{layer_i}.mlp.up_proj.weight": loaded[f"layers.{layer_i}.feed_forward.w3.weight"],
																		    f"model.layers.{layer_i}.input_layernorm.weight": loaded[f"layers.{layer_i}.attention_norm.weight"],
																		    f"model.layers.{layer_i}.post_attention_layernorm.weight": loaded[f"layers.{layer_i}.ffn_norm.weight"],
																		}
												else:
																		# Sharded
																		# Note that attention.w{q,k,v,o}, feed_fordward.w[1,2,3], attention_norm.weight and ffn_norm.weight share
																		# the same storage object, saving attention_norm and ffn_norm will save other weights too, which is
																		# redundant as other weights will be stitched from multiple shards. To avoid that, they are cloned.

																		A__					:       Any     ={
																		    f"model.layers.{layer_i}.input_layernorm.weight": loaded[0][
																		        f"layers.{layer_i}.attention_norm.weight"
																		    ].clone(),
																		    f"model.layers.{layer_i}.post_attention_layernorm.weight": loaded[0][
																		        f"layers.{layer_i}.ffn_norm.weight"
																		    ].clone(),
																		}
																		A__					:       Optional[Any]     =permute(
																		    torch.cat(
																		        [
																		            loaded[i][f"layers.{layer_i}.attention.wq.weight"].view(__snake_case,		__snake_case,		__snake_case	)
																		            for i in range(__snake_case	)
																		        ],		dim=0,		).reshape(__snake_case,		__snake_case	)	)
																		A__					:       int     =permute(
																		    torch.cat(
																		        [
																		            loaded[i][f"layers.{layer_i}.attention.wk.weight"].view(
																		                __snake_case,		__snake_case,		__snake_case	)
																		            for i in range(__snake_case	)
																		        ],		dim=0,		).reshape(__snake_case,		__snake_case	),		__snake_case,		__snake_case,		__snake_case,		)
																		A__					:       int     =torch.cat(
																		    [
																		        loaded[i][f"layers.{layer_i}.attention.wv.weight"].view(
																		            __snake_case,		__snake_case,		__snake_case	)
																		        for i in range(__snake_case	)
																		    ],		dim=0,		).reshape(__snake_case,		__snake_case	)

																		A__					:       List[str]     =torch.cat(
																		    [loaded[i][f"layers.{layer_i}.attention.wo.weight"] for i in range(__snake_case	)],		dim=1	)
																		A__					:       Optional[int]     =torch.cat(
																		    [loaded[i][f"layers.{layer_i}.feed_forward.w1.weight"] for i in range(__snake_case	)],		dim=0	)
																		A__					:       str     =torch.cat(
																		    [loaded[i][f"layers.{layer_i}.feed_forward.w2.weight"] for i in range(__snake_case	)],		dim=1	)
																		A__					:       List[str]     =torch.cat(
																		    [loaded[i][f"layers.{layer_i}.feed_forward.w3.weight"] for i in range(__snake_case	)],		dim=0	)

												A__					:       List[Any]     =inv_freq
												for k, v in state_dict.items():
																		A__					:       Optional[Any]     =filename
																		param_count += v.numel()
												torch.save(__snake_case,		os.path.join(__snake_case,		__snake_case	)	)

						A__					:       Tuple     =f"pytorch_model-{n_layers + 1}-of-{n_layers + 1}.bin"
						if model_size == "7B":
												# Unsharded
												A__					:       Tuple     ={
												    """model.embed_tokens.weight""": loaded["""tok_embeddings.weight"""],
												    """model.norm.weight""": loaded["""norm.weight"""],
												    """lm_head.weight""": loaded["""output.weight"""],
												}
						else:
												A__					:       Any     ={
												    """model.norm.weight""": loaded[0]["""norm.weight"""],
												    """model.embed_tokens.weight""": torch.cat(
												        [loaded[i]["""tok_embeddings.weight"""] for i in range(__snake_case	)],		dim=1	),
												    """lm_head.weight""": torch.cat([loaded[i]["""output.weight"""] for i in range(__snake_case	)],		dim=0	),
												}

						for k, v in state_dict.items():
												A__					:       int     =filename
												param_count += v.numel()
						torch.save(__snake_case,		os.path.join(__snake_case,		__snake_case	)	)

						# Write configs
						A__					:       Union[str, Any]     ={"""total_size""": param_count * 2}
						write_json(__snake_case,		os.path.join(__snake_case,		"""pytorch_model.bin.index.json"""	)	)
						A__					:       Optional[Any]     =params["""ffn_dim_multiplier"""] if """ffn_dim_multiplier""" in params else 1
						A__					:       List[Any]     =params["""multiple_of"""] if """multiple_of""" in params else 256
						A__					:       int     =LlamaConfig(
						    hidden_size=__snake_case,		intermediate_size=compute_intermediate_size(__snake_case,		__snake_case,		__snake_case	),		num_attention_heads=params["""n_heads"""],		num_hidden_layers=params["""n_layers"""],		rms_norm_eps=params["""norm_eps"""],		num_key_value_heads=__snake_case,		)
						config.save_pretrained(__snake_case	)

						# Make space so we can load the model properly now.
						del state_dict
						del loaded
						gc.collect()

						print("""Loading the checkpoint in a Llama model."""	)
						A__					:       List[Any]     =LlamaForCausalLM.from_pretrained(__snake_case,		torch_dtype=torch.floataa,		low_cpu_mem_usage=__snake_case	)
						# Avoid saving this as part of the config.
						del model.config._name_or_path

						print("""Saving in the Transformers format."""	)
						model.save_pretrained(__snake_case,		safe_serialization=__snake_case	)
						shutil.rmtree(__snake_case	)
def 						__lowerCamelCase	(       __snake_case					:		Union[str, Any],		__snake_case					:		Dict	)	->							Tuple:






						"""simple docstring"""





						A__					:       List[Any]     =LlamaTokenizer if LlamaTokenizerFast is None else LlamaTokenizerFast
						print(f"Saving a {tokenizer_class.__name__} to {tokenizer_path}."	)
						A__					:       List[str]     =tokenizer_class(__snake_case	)
						tokenizer.save_pretrained(__snake_case	)
def 						__lowerCamelCase	(       )	->							Any:






						"""simple docstring"""





						A__					:       List[str]     =argparse.ArgumentParser()
						parser.add_argument(
						    """--input_dir""",		help="""Location of LLaMA weights, which contains tokenizer.model and model folders""",		)
						parser.add_argument(
						    """--model_size""",		choices=["""7B""", """7Bf""", """13B""", """13Bf""", """30B""", """65B""", """70B""", """70Bf""", """tokenizer_only"""],		)
						parser.add_argument(
						    """--output_dir""",		help="""Location to write HF model and tokenizer""",		)
						parser.add_argument("""--safe_serialization""",		type=__snake_case,		help="""Whether or not to save using `safetensors`."""	)
						A__					:       Any     =parser.parse_args()
						if args.model_size != "tokenizer_only":
												write_model(
												    model_path=args.output_dir,		input_base_path=os.path.join(args.input_dir,		args.model_size	),		model_size=args.model_size,		safe_serialization=args.safe_serialization,		)
						A__					:       List[Any]     =os.path.join(args.input_dir,		"""tokenizer.model"""	)
						write_tokenizer(args.output_dir,		__snake_case	)


if __name__ == "__main__":
		main()

