






import gc
import unittest

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

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    EulerAncestralDiscreteScheduler,
    PNDMScheduler,
    StableDiffusionModelEditingPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu, skip_mps

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()







@skip_mps
class  lowerCamelCase_   (		a_       ,  a_       ,  a_       ,  unittest.TestCase   ):
							SCREAMING_SNAKE_CASE_     =	StableDiffusionModelEditingPipeline
							SCREAMING_SNAKE_CASE_     =	TEXT_TO_IMAGE_PARAMS
							SCREAMING_SNAKE_CASE_     =	TEXT_TO_IMAGE_BATCH_PARAMS
							SCREAMING_SNAKE_CASE_     =	TEXT_TO_IMAGE_IMAGE_PARAMS
							SCREAMING_SNAKE_CASE_     =	TEXT_TO_IMAGE_IMAGE_PARAMS
							def        SCREAMING_SNAKE_CASE_       (  self						:    Dict    ):





								'''simple docstring'''



								torch.manual_seed(0    )
								a				=	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				=	DDIMScheduler()
								torch.manual_seed(0    )
								a				=	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							,)
								torch.manual_seed(0    )
								a				=	CLIPTextConfig(
								    bos_token_id=0							,eos_token_id=2							,hidden_size=32							,intermediate_size=37							,layer_norm_eps=1e-05							,num_attention_heads=4							,num_hidden_layers=5							,pad_token_id=1							,vocab_size=10_00							,)
								a				=	CLIPTextModel(lowercase__    )
								a				=	CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''    )

								a				=	{
								    '''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						:    List[str]							,__lowerCamelCase						:    List[Any]							,__lowerCamelCase						:    Any=0    ):





								'''simple docstring'''



								a				=	torch.manual_seed(lowercase__    )
								a				=	{
								    '''prompt''': '''A field of roses''',
								    '''generator''': generator,
								    # Setting height and width to None to prevent OOMs on CPU.
								    '''height''': None,
								    '''width''': None,
								    '''num_inference_steps''': 2,
								    '''guidance_scale''': 6.0,
								    '''output_type''': '''numpy''',
								}
								return inputs
							def        SCREAMING_SNAKE_CASE_       (  self						:    Dict    ):





								'''simple docstring'''



								a				=	'''cpu'''  # ensure determinism for the device-dependent torch.Generator
								a				=	self.get_dummy_components()
								a				=	StableDiffusionModelEditingPipeline(**lowercase__    )
								a				=	sd_pipe.to(lowercase__    )
								sd_pipe.set_progress_bar_config(disable=lowercase__    )

								a				=	self.get_dummy_inputs(lowercase__    )
								a				=	sd_pipe(**lowercase__    ).images
								a				=	image[0, -3:, -3:, -1]
								assert image.shape == (1, 64, 64, 3)

								a				=	np.array([0.4_755, 0.5_132, 0.4_976, 0.3_904, 0.3_554, 0.4_765, 0.5_139, 0.5_158, 0.4_889]    )

								assert np.abs(image_slice.flatten() - expected_slice    ).max() < 1e-2
							def        SCREAMING_SNAKE_CASE_       (  self						:    List[str]    ):





								'''simple docstring'''



								a				=	'''cpu'''  # ensure determinism for the device-dependent torch.Generator
								a				=	self.get_dummy_components()
								a				=	StableDiffusionModelEditingPipeline(**lowercase__    )
								a				=	sd_pipe.to(lowercase__    )
								sd_pipe.set_progress_bar_config(disable=lowercase__    )

								a				=	self.get_dummy_inputs(lowercase__    )
								a				=	'''french fries'''
								a				=	sd_pipe(**lowercase__							,negative_prompt=lowercase__    )
								a				=	output.images
								a				=	image[0, -3:, -3:, -1]

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

								a				=	np.array([0.4_992, 0.5_101, 0.5_004, 0.3_949, 0.3_604, 0.4_735, 0.5_216, 0.5_204, 0.4_913]    )

								assert np.abs(image_slice.flatten() - expected_slice    ).max() < 1e-2
							def        SCREAMING_SNAKE_CASE_       (  self						:    Tuple    ):





								'''simple docstring'''



								a				=	'''cpu'''  # ensure determinism for the device-dependent torch.Generator
								a				=	self.get_dummy_components()
								a				=	EulerAncestralDiscreteScheduler(
								    beta_start=0.00_085							,beta_end=0.012							,beta_schedule='''scaled_linear'''    )
								a				=	StableDiffusionModelEditingPipeline(**lowercase__    )
								a				=	sd_pipe.to(lowercase__    )
								sd_pipe.set_progress_bar_config(disable=lowercase__    )

								a				=	self.get_dummy_inputs(lowercase__    )
								a				=	sd_pipe(**lowercase__    ).images
								a				=	image[0, -3:, -3:, -1]

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

								a				=	np.array([0.4_747, 0.5_372, 0.4_779, 0.4_982, 0.5_543, 0.4_816, 0.5_238, 0.4_904, 0.5_027]    )

								assert np.abs(image_slice.flatten() - expected_slice    ).max() < 1e-2
							def        SCREAMING_SNAKE_CASE_       (  self						:    Union[str, Any]    ):





								'''simple docstring'''



								a				=	'''cpu'''  # ensure determinism for the device-dependent torch.Generator
								a				=	self.get_dummy_components()
								a				=	PNDMScheduler()
								a				=	StableDiffusionModelEditingPipeline(**lowercase__    )
								a				=	sd_pipe.to(lowercase__    )
								sd_pipe.set_progress_bar_config(disable=lowercase__    )

								a				=	self.get_dummy_inputs(lowercase__    )
								# the pipeline does not expect pndm so test if it raises error.
								with self.assertRaises(lowercase__    ):
									a				=	sd_pipe(**lowercase__    ).images
							def        SCREAMING_SNAKE_CASE_       (  self						:    Optional[Any]    ):





								'''simple docstring'''



								super().test_inference_batch_single_identical(expected_max_diff=5e-3    )




							def        SCREAMING_SNAKE_CASE_       (  self						:    Tuple    ):





								'''simple docstring'''



								super().test_attention_slicing_forward_pass(expected_max_diff=5e-3    )







@slow
@require_torch_gpu
class  lowerCamelCase_   (		unittest.TestCase   ):
							def        SCREAMING_SNAKE_CASE_       (  self						:    Optional[Any]    ):





								'''simple docstring'''



								super().tearDown()
								gc.collect()
								torch.cuda.empty_cache()
							def        SCREAMING_SNAKE_CASE_       (  self						:    List[str]							,__lowerCamelCase						:    List[str]=0    ):





								'''simple docstring'''



								a				=	torch.manual_seed(lowercase__    )
								a				=	{
								    '''prompt''': '''A field of roses''',
								    '''generator''': generator,
								    '''num_inference_steps''': 3,
								    '''guidance_scale''': 7.5,
								    '''output_type''': '''numpy''',
								}
								return inputs
							def        SCREAMING_SNAKE_CASE_       (  self						:    str    ):





								'''simple docstring'''



								a				=	'''CompVis/stable-diffusion-v1-4'''
								a				=	StableDiffusionModelEditingPipeline.from_pretrained(lowercase__							,safety_checker=lowercase__    )
								pipe.to(lowercase__    )
								pipe.set_progress_bar_config(disable=lowercase__    )
								pipe.enable_attention_slicing()

								a				=	self.get_inputs()
								a				=	pipe(**lowercase__    ).images
								a				=	image[0, -3:, -3:, -1].flatten()

								assert image.shape == (1, 5_12, 5_12, 3)

								a				=	np.array(
								    [0.6_749_496, 0.6_386_453, 0.51_443_267, 0.66_094_905, 0.61_921_215, 0.5_491_332, 0.5_744_417, 0.58_075_106, 0.5_174_658]    )

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

								# make sure image changes after editing
								pipe.edit_model('''A pack of roses'''							,'''A pack of blue roses'''    )

								a				=	pipe(**lowercase__    ).images
								a				=	image[0, -3:, -3:, -1].flatten()

								assert image.shape == (1, 5_12, 5_12, 3)

								assert np.abs(expected_slice - image_slice    ).max() > 1e-1




							def        SCREAMING_SNAKE_CASE_       (  self						:    int    ):





								'''simple docstring'''



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

								a				=	'''CompVis/stable-diffusion-v1-4'''
								a				=	DDIMScheduler.from_pretrained(lowercase__							,subfolder='''scheduler'''    )
								a				=	StableDiffusionModelEditingPipeline.from_pretrained(
								    lowercase__							,scheduler=lowercase__							,safety_checker=lowercase__    )
								a				=	pipe.to(lowercase__    )
								pipe.set_progress_bar_config(disable=lowercase__    )
								pipe.enable_attention_slicing(1    )
								pipe.enable_sequential_cpu_offload()

								a				=	self.get_inputs()
								a				=	pipe(**lowercase__    )

								a				=	torch.cuda.max_memory_allocated()
								# make sure that less than 4.4 GB is allocated
								assert mem_bytes < 4.4 * 10**9






import argparse

import torch
from flax.training.checkpoints import restore_checkpoint

from transformers import FNetConfig, FNetForPreTraining
from transformers.utils import logging


logging.set_verbosity_info()






def       __lowerCamelCase		(						lowerCAmelCase__	,  lowerCAmelCase__	,  lowerCAmelCase__	):
       # Initialise PyTorch model
       lowerCAmelCase__     	=					FNetConfig.from_json_file(lowerCAmelCase__	)
       print(F"""Building PyTorch model from configuration: {config}"""	)
       lowerCAmelCase__     	=					FNetForPreTraining(lowerCAmelCase__	)

       lowerCAmelCase__     	=					restore_checkpoint(lowerCAmelCase__	,  lowerCAmelCase__	)
       lowerCAmelCase__     	=					checkpoint_dict['target']

       # Embeddings
       # Position IDs
       lowerCAmelCase__     	=					fnet_pretraining_model.state_dict()

       lowerCAmelCase__     	=					state_dict['fnet.embeddings.position_ids']
       lowerCAmelCase__     	=					{'fnet.embeddings.position_ids': position_ids}
       # Embedding Layers
       lowerCAmelCase__     	=					torch.tensor(
           pretrained_model_params['encoder']['embedder']['word']['embedding']	)
       lowerCAmelCase__     	=					torch.tensor(
           pretrained_model_params['encoder']['embedder']['position']['embedding'][0]	)
       lowerCAmelCase__     	=					torch.tensor(
           pretrained_model_params['encoder']['embedder']['type']['embedding']	)
       lowerCAmelCase__     	=					torch.tensor(
           pretrained_model_params['encoder']['embedder']['hidden_mapping_in']['kernel']	).T
       lowerCAmelCase__     	=					torch.tensor(
           pretrained_model_params['encoder']['embedder']['hidden_mapping_in']['bias']	)
       lowerCAmelCase__     	=					torch.tensor(
           pretrained_model_params['encoder']['embedder']['layer_norm']['scale']	)
       lowerCAmelCase__     	=					torch.tensor(
           pretrained_model_params['encoder']['embedder']['layer_norm']['bias']	)

       # Encoder Layers
       for layer in range(config.num_hidden_layers	):
              lowerCAmelCase__     	=					torch.tensor(
                  pretrained_model_params['encoder'][F"""encoder_{layer}"""]['mixing_layer_norm']['scale']	)
              lowerCAmelCase__     	=					torch.tensor(
                  pretrained_model_params['encoder'][F"""encoder_{layer}"""]['mixing_layer_norm']['bias']	)

              lowerCAmelCase__     	=					torch.tensor(
                  pretrained_model_params['encoder'][F"""feed_forward_{layer}"""]['intermediate']['kernel']	).T
              lowerCAmelCase__     	=					torch.tensor(
                  pretrained_model_params['encoder'][F"""feed_forward_{layer}"""]['intermediate']['bias']	)

              lowerCAmelCase__     	=					torch.tensor(
                  pretrained_model_params['encoder'][F"""feed_forward_{layer}"""]['output']['kernel']	).T
              lowerCAmelCase__     	=					torch.tensor(
                  pretrained_model_params['encoder'][F"""feed_forward_{layer}"""]['output']['bias']	)

              lowerCAmelCase__     	=					torch.tensor(
                  pretrained_model_params['encoder'][F"""encoder_{layer}"""]['output_layer_norm']['scale']	)
              lowerCAmelCase__     	=					torch.tensor(
                  pretrained_model_params['encoder'][F"""encoder_{layer}"""]['output_layer_norm']['bias']	)

       # Pooler Layers
       lowerCAmelCase__     	=					torch.tensor(pretrained_model_params['encoder']['pooler']['kernel']	).T
       lowerCAmelCase__     	=					torch.tensor(pretrained_model_params['encoder']['pooler']['bias']	)

       # Masked LM Layers
       lowerCAmelCase__     	=					torch.tensor(
           pretrained_model_params['predictions_dense']['kernel']	).T
       lowerCAmelCase__     	=					torch.tensor(
           pretrained_model_params['predictions_dense']['bias']	)
       lowerCAmelCase__     	=					torch.tensor(
           pretrained_model_params['predictions_layer_norm']['scale']	)
       lowerCAmelCase__     	=					torch.tensor(
           pretrained_model_params['predictions_layer_norm']['bias']	)
       lowerCAmelCase__     	=					torch.tensor(
           pretrained_model_params['encoder']['embedder']['word']['embedding']	)
       lowerCAmelCase__     	=					torch.tensor(
           pretrained_model_params['predictions_output']['output_bias']	)
       lowerCAmelCase__     	=					torch.tensor(pretrained_model_params['predictions_output']['output_bias']	)

       # Seq Relationship Layers
       lowerCAmelCase__     	=					torch.tensor(
           pretrained_model_params['classification']['output_kernel']	)
       lowerCAmelCase__     	=					torch.tensor(
           pretrained_model_params['classification']['output_bias']	)

       # Load State Dict
       fnet_pretraining_model.load_state_dict(lowerCAmelCase__	)

       # Save PreTrained
       print(F"""Saving pretrained model to {save_path}"""	)
       fnet_pretraining_model.save_pretrained(lowerCAmelCase__	)


if __name__ == "__main__":
    lowerCAmelCase__							= argparse.ArgumentParser()
    # Required parameters
    parser.add_argument(
        '--flax_checkpoint_path', default=None, type=str, required=True, help='Path to the TensorFlow checkpoint path.'
    )
    parser.add_argument(
        '--fnet_config_file',
        default=None,
        type=str,
        required=True,
        help=(
            'The config json file corresponding to the pre-trained FNet model. \n'
            'This specifies the model architecture.'
        ),
    )
    parser.add_argument('--save_path', default=None, type=str, required=True, help='Path to the output model.')
    lowerCAmelCase__							= parser.parse_args()
    convert_flax_checkpoint_to_pytorch(args.flax_checkpoint_path, args.fnet_config_file, args.save_path)
