


import gc
import random
import unittest

import torch

from diffusers import (
    IFImgaImgPipeline,
    IFImgaImgSuperResolutionPipeline,
    IFInpaintingPipeline,
    IFInpaintingSuperResolutionPipeline,
    IFPipeline,
    IFSuperResolutionPipeline,
)
from diffusers.models.attention_processor import AttnAddedKVProcessor
from diffusers.utils.import_utils import is_xformers_available
from diffusers.utils.testing_utils import floats_tensor, load_numpy, require_torch_gpu, skip_mps, slow, torch_device

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineTesterMixin, assert_mean_pixel_difference
from . import IFPipelineTesterMixin



@skip_mps
class 				A       (			UpperCAmelCase__					,    UpperCAmelCase__					,    unittest.TestCase				):
					'''simple docstring'''






					A__     =					IFPipeline
					A__     =					TEXT_TO_IMAGE_PARAMS - {'''width''', '''height''', '''latents'''}
					A__     =					TEXT_TO_IMAGE_BATCH_PARAMS
					A__     =					PipelineTesterMixin.required_optional_params - {'''latents'''}

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




							"""simple docstring"""


							return self._get_dummy_components()

					def     lowerCamelCase__	(self       :		int ,   _UpperCAmelCase       :		List[Any] ,   _UpperCAmelCase       :		str=0						)		->  Optional[Any]:




							"""simple docstring"""


							if str(_UpperCAmelCase						).startswith("""mps"""						):
									lowercase__	     =						torch.manual_seed(_UpperCAmelCase						)
							else:
									lowercase__	     =						torch.Generator(device=_UpperCAmelCase						).manual_seed(_UpperCAmelCase						)

							lowercase__	     =						{
							    """prompt""": """A painting of a squirrel eating a burger""",
							    """generator""": generator,
							    """num_inference_steps""": 2,
							    """output_type""": """numpy""",
							}

							return inputs

					def     lowerCamelCase__	(self       :		int						)		->  List[Any]:




							"""simple docstring"""


							self._test_save_load_optional_components()

					@unittest.skipIf(torch_device != """cuda""" ,   reason="""float16 requires CUDA"""						)
					def     lowerCamelCase__	(self       :		Dict						)		->  List[Any]:




							"""simple docstring"""


							super().test_save_load_floataa(expected_max_diff=1E-1						)

					def     lowerCamelCase__	(self       :		str						)		->  int:




							"""simple docstring"""


							self._test_attention_slicing_forward_pass(expected_max_diff=1E-2						)

					def     lowerCamelCase__	(self       :		int						)		->  Any:




							"""simple docstring"""


							self._test_save_load_local()

					def     lowerCamelCase__	(self       :		str						)		->  Optional[Any]:




							"""simple docstring"""


							self._test_inference_batch_single_identical(
							    expected_max_diff=1E-2 ,   )


					@unittest.skipIf(
					    torch_device != """cuda""" or not is_xformers_available() ,   reason="""XFormers attention is only available with CUDA and `xformers` installed""" ,   )
					def     lowerCamelCase__	(self       :		Optional[int]						)		->  Optional[int]:




							"""simple docstring"""


							self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=1E-3						)






@slow
@require_torch_gpu
class 				A       (			unittest.TestCase				):
					'''simple docstring'''






					def     lowerCamelCase__	(self       :		Tuple						)		->  str:




							"""simple docstring"""


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

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




							"""simple docstring"""


							lowercase__	     =						IFPipeline.from_pretrained("""DeepFloyd/IF-I-XL-v1.0""" ,   variant="""fp16""" ,   torch_dtype=torch.floataa						)

							lowercase__	     =						IFSuperResolutionPipeline.from_pretrained(
							    """DeepFloyd/IF-II-L-v1.0""" ,   variant="""fp16""" ,   torch_dtype=torch.floataa ,   text_encoder=_UpperCAmelCase ,   tokenizer=_UpperCAmelCase						)

							# pre compute text embeddings and remove T5 to save memory

							pipe_a.text_encoder.to("""cuda"""						)

							lowercase__     ,      lowercase__	     =						pipe_a.encode_prompt("""anime turtle""" ,   device="""cuda"""						)

							del pipe_a.tokenizer
							del pipe_a.text_encoder
							gc.collect()

							lowercase__	     =						None
							lowercase__	     =						None

							pipe_a.enable_model_cpu_offload()
							pipe_a.enable_model_cpu_offload()

							pipe_a.unet.set_attn_processor(AttnAddedKVProcessor()						)
							pipe_a.unet.set_attn_processor(AttnAddedKVProcessor()						)

							self._test_if(_UpperCAmelCase ,   _UpperCAmelCase ,   _UpperCAmelCase ,   _UpperCAmelCase						)

							pipe_a.remove_all_hooks()
							pipe_a.remove_all_hooks()

							# img2img

							lowercase__	     =						IFImgaImgPipeline(**pipe_a.components						)
							lowercase__	     =						IFImgaImgSuperResolutionPipeline(**pipe_a.components						)

							pipe_a.enable_model_cpu_offload()
							pipe_a.enable_model_cpu_offload()

							pipe_a.unet.set_attn_processor(AttnAddedKVProcessor()						)
							pipe_a.unet.set_attn_processor(AttnAddedKVProcessor()						)

							self._test_if_imgaimg(_UpperCAmelCase ,   _UpperCAmelCase ,   _UpperCAmelCase ,   _UpperCAmelCase						)

							pipe_a.remove_all_hooks()
							pipe_a.remove_all_hooks()

							# inpainting

							lowercase__	     =						IFInpaintingPipeline(**pipe_a.components						)
							lowercase__	     =						IFInpaintingSuperResolutionPipeline(**pipe_a.components						)

							pipe_a.enable_model_cpu_offload()
							pipe_a.enable_model_cpu_offload()

							pipe_a.unet.set_attn_processor(AttnAddedKVProcessor()						)
							pipe_a.unet.set_attn_processor(AttnAddedKVProcessor()						)

							self._test_if_inpainting(_UpperCAmelCase ,   _UpperCAmelCase ,   _UpperCAmelCase ,   _UpperCAmelCase						)

					def     lowerCamelCase__	(self       :		Any ,   _UpperCAmelCase       :		Tuple ,   _UpperCAmelCase       :		str ,   _UpperCAmelCase       :		Dict ,   _UpperCAmelCase       :		List[str]						)		->  List[Any]:




							"""simple docstring"""


							_start_torch_memory_measurement()

							lowercase__	     =						torch.Generator(device="""cpu"""						).manual_seed(0						)
							lowercase__	     =						pipe_a(
							    prompt_embeds=_UpperCAmelCase ,   negative_prompt_embeds=_UpperCAmelCase ,   num_inference_steps=2 ,   generator=_UpperCAmelCase ,   output_type="""np""" ,   )

							lowercase__	     =						output.images[0]

							assert image.shape == (64, 64, 3)

							lowercase__	     =						torch.cuda.max_memory_allocated()
							assert mem_bytes < 13 * 10**9

							lowercase__	     =						load_numpy(
							    """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/if/test_if.npy"""						)
							assert_mean_pixel_difference(_UpperCAmelCase ,   _UpperCAmelCase						)

							# pipeline 2

							_start_torch_memory_measurement()

							lowercase__	     =						torch.Generator(device="""cpu"""						).manual_seed(0						)

							lowercase__	     =						floats_tensor((1, 3, 64, 64) ,   rng=random.Random(0						)						).to(_UpperCAmelCase						)

							lowercase__	     =						pipe_a(
							    prompt_embeds=_UpperCAmelCase ,   negative_prompt_embeds=_UpperCAmelCase ,   image=_UpperCAmelCase ,   generator=_UpperCAmelCase ,   num_inference_steps=2 ,   output_type="""np""" ,   )

							lowercase__	     =						output.images[0]

							assert image.shape == (256, 256, 3)

							lowercase__	     =						torch.cuda.max_memory_allocated()
							assert mem_bytes < 4 * 10**9

							lowercase__	     =						load_numpy(
							    """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/if/test_if_superresolution_stage_II.npy"""						)
							assert_mean_pixel_difference(_UpperCAmelCase ,   _UpperCAmelCase						)

					def     lowerCamelCase__	(self       :		Union[str, Any] ,   _UpperCAmelCase       :		Optional[Any] ,   _UpperCAmelCase       :		str ,   _UpperCAmelCase       :		str ,   _UpperCAmelCase       :		Optional[int]						)		->  Any:




							"""simple docstring"""


							_start_torch_memory_measurement()

							lowercase__	     =						floats_tensor((1, 3, 64, 64) ,   rng=random.Random(0						)						).to(_UpperCAmelCase						)

							lowercase__	     =						torch.Generator(device="""cpu"""						).manual_seed(0						)

							lowercase__	     =						pipe_a(
							    prompt_embeds=_UpperCAmelCase ,   negative_prompt_embeds=_UpperCAmelCase ,   image=_UpperCAmelCase ,   num_inference_steps=2 ,   generator=_UpperCAmelCase ,   output_type="""np""" ,   )

							lowercase__	     =						output.images[0]

							assert image.shape == (64, 64, 3)

							lowercase__	     =						torch.cuda.max_memory_allocated()
							assert mem_bytes < 10 * 10**9

							lowercase__	     =						load_numpy(
							    """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/if/test_if_img2img.npy"""						)
							assert_mean_pixel_difference(_UpperCAmelCase ,   _UpperCAmelCase						)

							# pipeline 2

							_start_torch_memory_measurement()

							lowercase__	     =						torch.Generator(device="""cpu"""						).manual_seed(0						)

							lowercase__	     =						floats_tensor((1, 3, 256, 256) ,   rng=random.Random(0						)						).to(_UpperCAmelCase						)
							lowercase__	     =						floats_tensor((1, 3, 64, 64) ,   rng=random.Random(0						)						).to(_UpperCAmelCase						)

							lowercase__	     =						pipe_a(
							    prompt_embeds=_UpperCAmelCase ,   negative_prompt_embeds=_UpperCAmelCase ,   image=_UpperCAmelCase ,   original_image=_UpperCAmelCase ,   generator=_UpperCAmelCase ,   num_inference_steps=2 ,   output_type="""np""" ,   )

							lowercase__	     =						output.images[0]

							assert image.shape == (256, 256, 3)

							lowercase__	     =						torch.cuda.max_memory_allocated()
							assert mem_bytes < 4 * 10**9

							lowercase__	     =						load_numpy(
							    """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/if/test_if_img2img_superresolution_stage_II.npy"""						)
							assert_mean_pixel_difference(_UpperCAmelCase ,   _UpperCAmelCase						)


					def     lowerCamelCase__	(self       :		Dict ,   _UpperCAmelCase       :		List[str] ,   _UpperCAmelCase       :		Union[str, Any] ,   _UpperCAmelCase       :		Union[str, Any] ,   _UpperCAmelCase       :		str						)		->  Dict:




							"""simple docstring"""


							_start_torch_memory_measurement()

							lowercase__	     =						floats_tensor((1, 3, 64, 64) ,   rng=random.Random(0						)						).to(_UpperCAmelCase						)
							lowercase__	     =						floats_tensor((1, 3, 64, 64) ,   rng=random.Random(1						)						).to(_UpperCAmelCase						)

							lowercase__	     =						torch.Generator(device="""cpu"""						).manual_seed(0						)
							lowercase__	     =						pipe_a(
							    prompt_embeds=_UpperCAmelCase ,   negative_prompt_embeds=_UpperCAmelCase ,   image=_UpperCAmelCase ,   mask_image=_UpperCAmelCase ,   num_inference_steps=2 ,   generator=_UpperCAmelCase ,   output_type="""np""" ,   )

							lowercase__	     =						output.images[0]

							assert image.shape == (64, 64, 3)

							lowercase__	     =						torch.cuda.max_memory_allocated()
							assert mem_bytes < 10 * 10**9

							lowercase__	     =						load_numpy(
							    """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/if/test_if_inpainting.npy"""						)
							assert_mean_pixel_difference(_UpperCAmelCase ,   _UpperCAmelCase						)

							# pipeline 2

							_start_torch_memory_measurement()

							lowercase__	     =						torch.Generator(device="""cpu"""						).manual_seed(0						)

							lowercase__	     =						floats_tensor((1, 3, 64, 64) ,   rng=random.Random(0						)						).to(_UpperCAmelCase						)
							lowercase__	     =						floats_tensor((1, 3, 256, 256) ,   rng=random.Random(0						)						).to(_UpperCAmelCase						)
							lowercase__	     =						floats_tensor((1, 3, 256, 256) ,   rng=random.Random(1						)						).to(_UpperCAmelCase						)

							lowercase__	     =						pipe_a(
							    prompt_embeds=_UpperCAmelCase ,   negative_prompt_embeds=_UpperCAmelCase ,   image=_UpperCAmelCase ,   mask_image=_UpperCAmelCase ,   original_image=_UpperCAmelCase ,   generator=_UpperCAmelCase ,   num_inference_steps=2 ,   output_type="""np""" ,   )

							lowercase__	     =						output.images[0]

							assert image.shape == (256, 256, 3)

							lowercase__	     =						torch.cuda.max_memory_allocated()
							assert mem_bytes < 4 * 10**9

							lowercase__	     =						load_numpy(
							    """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/if/test_if_inpainting_superresolution_stage_II.npy"""						)
							assert_mean_pixel_difference(_UpperCAmelCase ,   _UpperCAmelCase						)






def 					UpperCamelCase   (      )  ->    Any:
		"""simple docstring"""




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






import argparse
import json
import os

import torch
from torch import nn

from transformers import NllbMoeConfig, NllbMoeModel
from transformers.modeling_utils import dtype_byte_size
from transformers.utils import WEIGHTS_INDEX_NAME, WEIGHTS_NAME






def     SCREAMING_SNAKE_CASE(      __UpperCamelCase )	->						List[Any]:
				a__	:							Any      						= [
				    "encoder.version",
				    "decoder.version",
				    "model.encoder.version",
				    "model.decoder.version",
				    "decoder.output_projection.weight",
				    "_float_tensor",
				    "encoder.embed_positions._float_tensor",
				    "decoder.embed_positions._float_tensor",
				]
				for k in ignore_keys:
								state_dict.pop(__UpperCamelCase	,				__UpperCamelCase )






def     SCREAMING_SNAKE_CASE(      __UpperCamelCase )	->						Optional[Any]:
				a__						,   a__	:							Optional[int]      						= emb.weight.shape
				a__	:							Union[str, Any]      						= nn.Linear(__UpperCamelCase	,				__UpperCamelCase	,				bias=__UpperCamelCase )
				a__	:							str      						= emb.weight.data
				return lin_layer






def     SCREAMING_SNAKE_CASE(      __UpperCamelCase	,				__UpperCamelCase=None )	->						Optional[int]:
				a__	:							Union[str, Any]      						= {}
				for old_key in state_dict.keys():
								a__	:							List[Any]      						= old_key
								if "moe_layer.experts." in key:
												if expert_idx is not None:
																a__	:							List[str]      						= key.replace("moe_layer.experts.0"	,				F'ffn.experts.expert_{expert_idx}' )
												else:
																a__	:							Optional[Any]      						= key.replace("moe_layer.experts."	,				"ffn.experts.expert_" )
								if "gate" in key:
												a__	:							Optional[Any]      						= key.replace(".moe_layer.gate.wg"	,				".ffn.router.classifier" )
								if "fc2" and "experts" not in key:
												a__	:							int      						= key.replace(".fc2."	,				".ffn.fc2." )
								if "fc1" and "experts" not in key:
												a__	:							Any      						= key.replace(".fc1."	,				".ffn.fc1." )
								if ".encoder_attn." in key:
												a__	:							str      						= key.replace(".encoder_attn."	,				".cross_attention." )
								if "encoder_attn_layer_norm" in key:
												a__	:							Optional[int]      						= key.replace("encoder_attn_layer_norm"	,				"cross_attention_layer_norm" )
								if "final_layer_norm" in key:
												a__	:							int      						= key.replace("final_layer_norm"	,				"ff_layer_norm" )
								a__	:							Optional[int]      						= state_dict[old_key]
				return new_dict






def     SCREAMING_SNAKE_CASE(      __UpperCamelCase	,				__UpperCamelCase	,				__UpperCamelCase	,				__UpperCamelCase	,				__UpperCamelCase = WEIGHTS_NAME )	->						Dict:
				a__	:							Any      						= []
				a__	:							List[str]      						= 0
				os.makedirs(__UpperCamelCase	,				exist_ok=__UpperCamelCase )

				for expert in range(__UpperCamelCase ):
								a__	:							Union[str, Any]      						= switch_checkpoint_path + F'-rank-{expert}.pt'
								if os.path.isfile(__UpperCamelCase ):
												a__	:							str      						= torch.load(__UpperCamelCase )["model"]
												remove_ignore_keys_(__UpperCamelCase )
												a__	:							Tuple      						= rename_fairseq_keys(__UpperCamelCase	,				__UpperCamelCase )
												a__	:							str      						= os.path.join(
												    __UpperCamelCase	,				weights_name.replace(".bin"	,				F'-{len(__UpperCamelCase )+1:05d}-of-???.bin' ) )
												torch.save(__UpperCamelCase	,				__UpperCamelCase )
												sharded_state_dicts.append(expert_state.keys() )
												total_size += sum([value.numel() for key, value in expert_state.items()] ) * dtype_byte_size(
												    expert_state[list(__UpperCamelCase )[0]].dtype )

    # Add the last block
				a__	:							int      						= os.path.join(__UpperCamelCase	,				weights_name.replace(".bin"	,				F'-{len(__UpperCamelCase )+1:05d}-of-???.bin' ) )
				a__	:							Any      						= torch.load(switch_checkpoint_path + "-shared.pt" )["model"]
				remove_ignore_keys_(__UpperCamelCase )
				a__	:							Optional[Any]      						= rename_fairseq_keys(__UpperCamelCase	,				__UpperCamelCase )
				a__	:							List[Any]      						= shared_weights["decoder.embed_tokens.weight"]
				sharded_state_dicts.append(shared_weights.keys() )

				# If we only have the shared weights (dummy model/experts saved on the same file)
				if len(__UpperCamelCase ) == 1:
								a__	:							Union[str, Any]      						= os.path.join(__UpperCamelCase	,				__UpperCamelCase )
								torch.save(__UpperCamelCase	,				__UpperCamelCase )
								return {weights_name: sharded_state_dicts[0]}, None
				else:
								torch.save(__UpperCamelCase	,				__UpperCamelCase )
				# Otherwise, let's build the index
				a__	:							Any      						= {}
				for idx, shard in enumerate(__UpperCamelCase ):
								a__	:							Any      						= weights_name.replace(".bin"	,				F'-{idx+1:05d}-of-{len(__UpperCamelCase ):05d}.bin' )
								a__	:							Optional[Any]      						= os.path.join(__UpperCamelCase	,				weights_name.replace(".bin"	,				F'-{idx+1:05d}-of-???.bin' ) )
								os.rename(__UpperCamelCase	,				os.path.join(__UpperCamelCase	,				__UpperCamelCase ) )
								for key in shard:
												a__	:							Any      						= shard_file

    # Add the metadata
				a__	:							Optional[int]      						= {"total_size": total_size}
				a__	:							Optional[Any]      						= {"metadata": metadata, "weight_map": weight_map}

				with open(os.path.join(__UpperCamelCase	,				__UpperCamelCase )	,				"w"	,				encoding="utf-8" ) as f:
								a__	:							Any      						= json.dumps(__UpperCamelCase	,				indent=2	,				sort_keys=__UpperCamelCase ) + "\n"
								f.write(__UpperCamelCase )

				return metadata, index


if __name__ == "__main__":
						lowerCamelCase		       =      argparse.ArgumentParser()
						# Required parameters
						parser.add_argument(
						    """--nllb_moe_checkpoint_path""",
						    default="""/home/arthur_huggingface_co/fairseq/weights/checkpoints/model_moe_54b/checkpoint_2_300000""",
						    type=str,
						    required=False,
						    help="""Path to a directory containing a folder per layer. Follows the original Google format.""",
						)
						parser.add_argument("""--dtype""", default="""float32""", type=str, required=False, help="""dtype of the saved model""")
						parser.add_argument(
						    """--pytorch_dump_folder_path""",
						    default="""/home/arthur_huggingface_co/fairseq/weights/checkpoints/hf-converted-moe-54b""",
						    type=str,
						    required=False,
						    help="""Path to the output pytorch model.""",
						)
						lowerCamelCase		       =      parser.parse_args()
						lowerCamelCase       ,   lowerCamelCase		       =      shard_on_the_fly(
						    args.nllb_moe_checkpoint_path,
						    args.pytorch_dump_folder_path,
						    1_28,
						    args.dtype,
						)

						lowerCamelCase		       =      NllbMoeConfig.from_pretrained(
						    """facebook/nllb-200-3.3B""", encoder_sparse_step=4, decoder_sparse_step=4, num_experts=1_28
						)
						config.save_pretrained(args.pytorch_dump_folder_path)
						lowerCamelCase		       =      NllbMoeModel.from_pretrained(args.pytorch_dump_folder_path)
						print("""Done""")
						model.save_pretrained(args.pytorch_dump_folder_path)


