






'''simple docstring'''




import argparse
import json
import os

import fairseq
import torch
from fairseq.data import Dictionary

# Register SEW's fairseq modules
from sew_asapp import tasks  # noqa: F401

from transformers import (
    SEWConfig,
    SEWForCTC,
    SEWModel,
    WavaVecaCTCTokenizer,
    WavaVecaFeatureExtractor,
    WavaVecaProcessor,
    logging,
)


logging.set_verbosity_info()
lowerCAmelCase__          =						logging.get_logger(__name__)

lowerCAmelCase__          =						{
    "post_extract_proj": "feature_projection",
    "encoder.pos_conv.0": "encoder.pos_conv_embed.conv",
    "self_attn.k_proj": "encoder.layers.*.attention.k_proj",
    "self_attn.v_proj": "encoder.layers.*.attention.v_proj",
    "self_attn.q_proj": "encoder.layers.*.attention.q_proj",
    "self_attn.out_proj": "encoder.layers.*.attention.out_proj",
    "self_attn_layer_norm": "encoder.layers.*.layer_norm",
    "fc1": "encoder.layers.*.feed_forward.intermediate_dense",
    "fc2": "encoder.layers.*.feed_forward.output_dense",
    "final_layer_norm": "encoder.layers.*.final_layer_norm",
    "encoder.upsample.0": "encoder.upsample.projection",
    "encoder.layer_norm": "encoder.layer_norm",
    "w2v_model.layer_norm": "layer_norm",
    "w2v_encoder.proj": "lm_head",
    "mask_emb": "masked_spec_embed",
}







def       __UpperCAmelCase						(  lowerCamelCase_     ,      lowerCamelCase_     ,      lowerCamelCase_     ,      lowerCamelCase_     ,      lowerCamelCase_)							->     Optional[Any]:
  for attribute in key.split('.'):
    UpperCamelCase__				:					int					=					getattr(__lowerCAmelCase     ,      __lowerCAmelCase)

  if weight_type is not None:
    UpperCamelCase__				:					int					=					getattr(__lowerCAmelCase     ,      __lowerCAmelCase).shape
  else:
    UpperCamelCase__				:					Optional[Any]					=					hf_pointer.shape

  assert hf_shape == value.shape, (
      f'Shape of hf {key + "." + weight_type if weight_type is not None else ""} is {hf_shape}, but should be'
      f' {value.shape} for {full_name}'
  )

  if weight_type == "weight":
    UpperCamelCase__				:					Dict					=					value
  elif weight_type == "weight_g":
    UpperCamelCase__				:					Dict					=					value
  elif weight_type == "weight_v":
    UpperCamelCase__				:					Optional[Any]					=					value
  elif weight_type == "bias":
    UpperCamelCase__				:					int					=					value
  else:
    UpperCamelCase__				:					str					=					value

  logger.info(f'{key + "." + weight_type if weight_type is not None else ""} was initialized from {full_name}.')







def       __UpperCAmelCase						(  lowerCamelCase_     ,      lowerCamelCase_     ,      lowerCamelCase_)							->     Dict:
  UpperCamelCase__				:					Optional[int]					=					[]
  UpperCamelCase__				:					Dict					=					fairseq_model.state_dict()

  UpperCamelCase__				:					str					=					hf_model.sew.feature_extractor if is_finetuned else hf_model.feature_extractor

  for name, value in fairseq_dict.items():
    UpperCamelCase__				:					List[str]					=					False
    if "conv_layers" in name:
      load_conv_layer(
          __lowerCAmelCase     ,      __lowerCAmelCase     ,      __lowerCAmelCase     ,      __lowerCAmelCase     ,      hf_model.config.feat_extract_norm == 'group'     ,      )
      UpperCamelCase__				:					List[Any]					=					True
    else:
      for key, mapped_key in MAPPING.items():
        UpperCamelCase__				:					Dict					=					"""sew.""" + mapped_key if (is_finetuned and mapped_key != """lm_head""") else mapped_key

        if key in name or key.split('w2v_model.')[-1] == name.split('.')[0]:
          UpperCamelCase__				:					Optional[Any]					=					True
          if "*" in mapped_key:
            UpperCamelCase__				:					List[Any]					=					name.split(__lowerCAmelCase)[0].split('.')[-2]
            UpperCamelCase__				:					List[str]					=					mapped_key.replace('*'     ,      __lowerCAmelCase)
          if "weight_g" in name:
            UpperCamelCase__				:					Union[str, Any]					=					"""weight_g"""
          elif "weight_v" in name:
            UpperCamelCase__				:					Union[str, Any]					=					"""weight_v"""
          elif "weight" in name:
            UpperCamelCase__				:					int					=					"""weight"""
          elif "bias" in name:
            UpperCamelCase__				:					List[Any]					=					"""bias"""
          else:
            UpperCamelCase__				:					Union[str, Any]					=					None
          set_recursively(__lowerCAmelCase     ,      __lowerCAmelCase     ,      __lowerCAmelCase     ,      __lowerCAmelCase     ,      __lowerCAmelCase)
        continue
    if not is_used:
      unused_weights.append(__lowerCAmelCase)

  logger.warning(f'Unused weights: {unused_weights}')







def       __UpperCAmelCase						(  lowerCamelCase_     ,      lowerCamelCase_     ,      lowerCamelCase_     ,      lowerCamelCase_     ,      lowerCamelCase_)							->     Optional[int]:
  UpperCamelCase__				:					str					=					full_name.split('conv_layers.')[-1]
  UpperCamelCase__				:					Union[str, Any]					=					name.split('.')
  UpperCamelCase__				:					Optional[int]					=					int(items[0])
  UpperCamelCase__				:					Any					=					int(items[1])

  if type_id == 0:
    if "bias" in name:
      assert value.shape == feature_extractor.conv_layers[layer_id].conv.bias.data.shape, (
          f'{full_name} has size {value.shape}, but'
          f' {feature_extractor.conv_layers[layer_id].conv.bias.data.shape} was found.'
      )
      UpperCamelCase__				:					Union[str, Any]					=					value
      logger.info(f'Feat extract conv layer {layer_id} was initialized from {full_name}.')
    elif "weight" in name:
      assert value.shape == feature_extractor.conv_layers[layer_id].conv.weight.data.shape, (
          f'{full_name} has size {value.shape}, but'
          f' {feature_extractor.conv_layers[layer_id].conv.weight.data.shape} was found.'
      )
      UpperCamelCase__				:					int					=					value
      logger.info(f'Feat extract conv layer {layer_id} was initialized from {full_name}.')
  elif (type_id == 2 and not use_group_norm) or (type_id == 2 and layer_id == 0 and use_group_norm):
    if "bias" in name:
      assert value.shape == feature_extractor.conv_layers[layer_id].layer_norm.bias.data.shape, (
          f'{full_name} has size {value.shape}, but {feature_extractor[layer_id].layer_norm.bias.data.shape} was'
          " found."
      )
      UpperCamelCase__				:					Optional[Any]					=					value
      logger.info(f'Feat extract layer norm weight of layer {layer_id} was initialized from {full_name}.')
    elif "weight" in name:
      assert value.shape == feature_extractor.conv_layers[layer_id].layer_norm.weight.data.shape, (
          f'{full_name} has size {value.shape}, but'
          f' {feature_extractor[layer_id].layer_norm.weight.data.shape} was found.'
      )
      UpperCamelCase__				:					Any					=					value
      logger.info(f'Feat extract layer norm weight of layer {layer_id} was initialized from {full_name}.')
  else:
    unused_weights.append(__lowerCAmelCase)







def       __UpperCAmelCase						(  lowerCamelCase_     ,      lowerCamelCase_)							->     List[str]:
  UpperCamelCase__				:					Union[str, Any]					=					SEWConfig()
  if is_finetuned:
    UpperCamelCase__				:					Tuple					=					model.wav_encoder.wav_model.cfg
  else:
    UpperCamelCase__				:					Dict					=					model.cfg

  UpperCamelCase__				:					Any					=					fs_config.conv_bias
  UpperCamelCase__				:					Optional[Any]					=					eval(fs_config.conv_feature_layers)
  UpperCamelCase__				:					Optional[Any]					=					[x[0] for x in conv_layers]
  UpperCamelCase__				:					List[Any]					=					[x[1] for x in conv_layers]
  UpperCamelCase__				:					Dict					=					[x[2] for x in conv_layers]
  UpperCamelCase__				:					List[str]					=					"""gelu"""
  UpperCamelCase__				:					str					=					"""layer""" if fs_config.extractor_mode == """layer_norm""" else """group"""
  UpperCamelCase__				:					Any					=					0.0
  UpperCamelCase__				:					Tuple					=					fs_config.activation_fn.name
  UpperCamelCase__				:					Tuple					=					fs_config.encoder_embed_dim
  UpperCamelCase__				:					Optional[Any]					=					0.02
  UpperCamelCase__				:					str					=					fs_config.encoder_ffn_embed_dim
  UpperCamelCase__				:					str					=					1e-5
  UpperCamelCase__				:					str					=					fs_config.encoder_layerdrop
  UpperCamelCase__				:					Optional[int]					=					fs_config.encoder_attention_heads
  UpperCamelCase__				:					int					=					fs_config.conv_pos_groups
  UpperCamelCase__				:					Optional[int]					=					fs_config.conv_pos
  UpperCamelCase__				:					Union[str, Any]					=					len(__lowerCAmelCase)
  UpperCamelCase__				:					Optional[int]					=					fs_config.encoder_layers
  UpperCamelCase__				:					Any					=					fs_config.squeeze_factor

  # take care of any params that are overridden by the Wav2VecCtc model
  if is_finetuned:
    UpperCamelCase__				:					Optional[Any]					=					model.cfg
    UpperCamelCase__				:					str					=					fs_config.final_dropout
    UpperCamelCase__				:					Union[str, Any]					=					fs_config.layerdrop
  UpperCamelCase__				:					List[str]					=					fs_config.activation_dropout
  UpperCamelCase__				:					Any					=					fs_config.mask_prob > 0 or fs_config.mask_channel_prob > 0
  UpperCamelCase__				:					str					=					fs_config.attention_dropout
  UpperCamelCase__				:					Optional[Any]					=					fs_config.dropout_input
  UpperCamelCase__				:					Union[str, Any]					=					fs_config.dropout
  UpperCamelCase__				:					Dict					=					fs_config.mask_channel_length
  UpperCamelCase__				:					int					=					fs_config.mask_channel_prob
  UpperCamelCase__				:					Tuple					=					fs_config.mask_length
  UpperCamelCase__				:					Tuple					=					fs_config.mask_prob

  UpperCamelCase__				:					Optional[int]					=					"""Wav2Vec2FeatureExtractor"""
  UpperCamelCase__				:					int					=					"""Wav2Vec2CTCTokenizer"""

  return config







@torch.no_grad()
def       __UpperCAmelCase						(  lowerCamelCase_     ,      lowerCamelCase_     ,      lowerCamelCase_=None     ,      lowerCamelCase_=None     ,      lowerCamelCase_=True)							->     str:

  if is_finetuned:
    UpperCamelCase__				:					List[str]					=					fairseq.checkpoint_utils.load_model_ensemble_and_task(
        [checkpoint_path]     ,      arg_overrides={'data': '/'.join(dict_path.split('/')[:-1])})
  else:
    UpperCamelCase__				:					Union[str, Any]					=					fairseq.checkpoint_utils.load_model_ensemble_and_task([checkpoint_path])

  if config_path is not None:
    UpperCamelCase__				:					str					=					SEWConfig.from_pretrained(__lowerCAmelCase)
  else:
    UpperCamelCase__				:					str					=					convert_config(model[0]     ,      __lowerCAmelCase)
  UpperCamelCase__				:					Optional[int]					=					model[0].eval()

  UpperCamelCase__				:					Dict					=					True if config.feat_extract_norm == """layer""" else False
  UpperCamelCase__				:					Union[str, Any]					=					WavaVecaFeatureExtractor(
      feature_size=1     ,      sampling_rate=16_000     ,      padding_value=0     ,      do_normalize=__lowerCAmelCase     ,      return_attention_mask=__lowerCAmelCase     ,      )

  if is_finetuned:
    if dict_path:
      UpperCamelCase__				:					str					=					Dictionary.load(__lowerCAmelCase)

      # important change bos & pad token id since CTC symbol is <pad> and
      # not <s> as in fairseq
      UpperCamelCase__				:					Tuple					=					target_dict.pad_index
      UpperCamelCase__				:					List[str]					=					target_dict.bos_index
      UpperCamelCase__				:					Any					=					target_dict.pad_index
      UpperCamelCase__				:					Any					=					target_dict.bos_index
      UpperCamelCase__				:					int					=					target_dict.eos_index
      UpperCamelCase__				:					List[str]					=					len(target_dict.symbols)
      UpperCamelCase__				:					Optional[int]					=					os.path.join(__lowerCAmelCase     ,      'vocab.json')
      if not os.path.isdir(__lowerCAmelCase):
        logger.error('--pytorch_dump_folder_path ({}) should be a directory'.format(__lowerCAmelCase))
        return
      os.makedirs(__lowerCAmelCase     ,      exist_ok=__lowerCAmelCase)
      with open(__lowerCAmelCase     ,      'w'     ,      encoding='utf-8') as vocab_handle:
        json.dump(target_dict.indices     ,      __lowerCAmelCase)
      UpperCamelCase__				:					Dict					=					WavaVecaCTCTokenizer(
          __lowerCAmelCase     ,      unk_token=target_dict.unk_word     ,      pad_token=target_dict.pad_word     ,      bos_token=target_dict.bos_word     ,      eos_token=target_dict.eos_word     ,      word_delimiter_token='|'     ,      do_lower_case=__lowerCAmelCase     ,      )
      UpperCamelCase__				:					int					=					WavaVecaProcessor(feature_extractor=__lowerCAmelCase     ,      tokenizer=__lowerCAmelCase)
      processor.save_pretrained(__lowerCAmelCase)

    UpperCamelCase__				:					Optional[Any]					=					SEWForCTC(__lowerCAmelCase)
  else:
    UpperCamelCase__				:					Dict					=					SEWModel(__lowerCAmelCase)
    feature_extractor.save_pretrained(__lowerCAmelCase)

  recursively_load_weights(__lowerCAmelCase     ,      __lowerCAmelCase     ,      __lowerCAmelCase)

  hf_model.save_pretrained(__lowerCAmelCase)


if __name__ == "__main__":
  lowerCAmelCase__          =						argparse.ArgumentParser()
  parser.add_argument('--pytorch_dump_folder_path', default=None, type=str, help='Path to the output PyTorch model.')
  parser.add_argument('--checkpoint_path', default=None, type=str, help='Path to fairseq checkpoint')
  parser.add_argument('--dict_path', default=None, type=str, help='Path to dict of fine-tuned model')
  parser.add_argument('--config_path', default=None, type=str, help='Path to hf config.json of model to convert')
  parser.add_argument(
      '--is_finetuned', action='store_true', help='Whether the model to convert is a fine-tuned model or not'
  )
  lowerCAmelCase__          =						parser.parse_args()
  convert_sew_checkpoint(
      args.checkpoint_path, args.pytorch_dump_folder_path, args.config_path, args.dict_path, args.is_finetuned
  )












'''simple docstring'''




import tempfile

import numpy as np
import torch
from transformers import AutoTokenizer, TaEncoderModel

from diffusers import DDPMScheduler, UNetaDConditionModel
from diffusers.models.attention_processor import AttnAddedKVProcessor
from diffusers.pipelines.deepfloyd_if import IFWatermarker
from diffusers.utils.testing_utils import torch_device

from ..test_pipelines_common import to_np







class       __lowercase					:

	def      __UpperCamelCase							(      self   :						Union[str, Any]):
			torch.manual_seed(0)
			UpperCamelCase__				:					Dict					=					TaEncoderModel.from_pretrained('hf-internal-testing/tiny-random-t5')

			torch.manual_seed(0)
			UpperCamelCase__				:					Union[str, Any]					=					AutoTokenizer.from_pretrained('hf-internal-testing/tiny-random-t5')

			torch.manual_seed(0)
			UpperCamelCase__				:					List[str]					=					UNetaDConditionModel(
			    sample_size=32  ,					layers_per_block=1  ,					block_out_channels=[32, 64]  ,					down_block_types=[
			        'ResnetDownsampleBlock2D',
			        'SimpleCrossAttnDownBlock2D',
			    ]  ,					mid_block_type='UNetMidBlock2DSimpleCrossAttn'  ,					up_block_types=['SimpleCrossAttnUpBlock2D', 'ResnetUpsampleBlock2D']  ,					in_channels=3  ,					out_channels=6  ,					cross_attention_dim=32  ,					encoder_hid_dim=32  ,					attention_head_dim=8  ,					addition_embed_type='text'  ,					addition_embed_type_num_heads=2  ,					cross_attention_norm='group_norm'  ,					resnet_time_scale_shift='scale_shift'  ,					act_fn='gelu'  ,					)
			unet.set_attn_processor(AttnAddedKVProcessor())  # For reproducibility tests

			torch.manual_seed(0)
			UpperCamelCase__				:					Optional[Any]					=					DDPMScheduler(
			    num_train_timesteps=1_000  ,					beta_schedule='squaredcos_cap_v2'  ,					beta_start=0.00_01  ,					beta_end=0.02  ,					thresholding=UpperCAmelCase_  ,					dynamic_thresholding_ratio=0.95  ,					sample_max_value=1.0  ,					prediction_type='epsilon'  ,					variance_type='learned_range'  ,					)

			torch.manual_seed(0)
			UpperCamelCase__				:					List[Any]					=					IFWatermarker()

			return {
			    "text_encoder": text_encoder,
			    "tokenizer": tokenizer,
			    "unet": unet,
			    "scheduler": scheduler,
			    "watermarker": watermarker,
			    "safety_checker": None,
			    "feature_extractor": None,
			}

	def      __UpperCamelCase							(      self   :						Dict):
			torch.manual_seed(0)
			UpperCamelCase__				:					List[Any]					=					TaEncoderModel.from_pretrained('hf-internal-testing/tiny-random-t5')

			torch.manual_seed(0)
			UpperCamelCase__				:					Any					=					AutoTokenizer.from_pretrained('hf-internal-testing/tiny-random-t5')

			torch.manual_seed(0)
			UpperCamelCase__				:					Any					=					UNetaDConditionModel(
			    sample_size=32  ,					layers_per_block=[1, 2]  ,					block_out_channels=[32, 64]  ,					down_block_types=[
			        'ResnetDownsampleBlock2D',
			        'SimpleCrossAttnDownBlock2D',
			    ]  ,					mid_block_type='UNetMidBlock2DSimpleCrossAttn'  ,					up_block_types=['SimpleCrossAttnUpBlock2D', 'ResnetUpsampleBlock2D']  ,					in_channels=6  ,					out_channels=6  ,					cross_attention_dim=32  ,					encoder_hid_dim=32  ,					attention_head_dim=8  ,					addition_embed_type='text'  ,					addition_embed_type_num_heads=2  ,					cross_attention_norm='group_norm'  ,					resnet_time_scale_shift='scale_shift'  ,					act_fn='gelu'  ,					class_embed_type='timestep'  ,					mid_block_scale_factor=1.4_14  ,					time_embedding_act_fn='gelu'  ,					time_embedding_dim=32  ,					)
			unet.set_attn_processor(AttnAddedKVProcessor())  # For reproducibility tests

			torch.manual_seed(0)
			UpperCamelCase__				:					str					=					DDPMScheduler(
			    num_train_timesteps=1_000  ,					beta_schedule='squaredcos_cap_v2'  ,					beta_start=0.00_01  ,					beta_end=0.02  ,					thresholding=UpperCAmelCase_  ,					dynamic_thresholding_ratio=0.95  ,					sample_max_value=1.0  ,					prediction_type='epsilon'  ,					variance_type='learned_range'  ,					)

			torch.manual_seed(0)
			UpperCamelCase__				:					List[str]					=					DDPMScheduler(
			    num_train_timesteps=1_000  ,					beta_schedule='squaredcos_cap_v2'  ,					beta_start=0.00_01  ,					beta_end=0.02  ,					)

			torch.manual_seed(0)
			UpperCamelCase__				:					Optional[Any]					=					IFWatermarker()

			return {
			    "text_encoder": text_encoder,
			    "tokenizer": tokenizer,
			    "unet": unet,
			    "scheduler": scheduler,
			    "image_noising_scheduler": image_noising_scheduler,
			    "watermarker": watermarker,
			    "safety_checker": None,
			    "feature_extractor": None,
			}

	def      __UpperCamelCase							(      self   :						Any):
			UpperCamelCase__				:					Dict					=					self.get_dummy_components()
			UpperCamelCase__				:					List[Any]					=					self.pipeline_class(**UpperCAmelCase_)
			pipe.to(UpperCAmelCase_)
			pipe.set_progress_bar_config(disable=UpperCAmelCase_)

			UpperCamelCase__				:					Tuple					=					self.get_dummy_inputs(UpperCAmelCase_)

			UpperCamelCase__				:					Optional[Any]					=					inputs['prompt']
			UpperCamelCase__				:					List[Any]					=					inputs['generator']
			UpperCamelCase__				:					Tuple					=					inputs['num_inference_steps']
			UpperCamelCase__				:					List[Any]					=					inputs['output_type']

			if "image" in inputs:
					UpperCamelCase__				:					Tuple					=					inputs['image']
			else:
					UpperCamelCase__				:					Union[str, Any]					=					None

			if "mask_image" in inputs:
					UpperCamelCase__				:					Optional[int]					=					inputs['mask_image']
			else:
					UpperCamelCase__				:					int					=					None

			if "original_image" in inputs:
					UpperCamelCase__				:					List[Any]					=					inputs['original_image']
			else:
					UpperCamelCase__				:					Optional[Any]					=					None

			UpperCamelCase__,	UpperCamelCase__				:					Any					=					pipe.encode_prompt(UpperCAmelCase_)

			# inputs with prompt converted to embeddings
			UpperCamelCase__				:					List[Any]					=					{
			    'prompt_embeds': prompt_embeds,
			    'negative_prompt_embeds': negative_prompt_embeds,
			    'generator': generator,
			    'num_inference_steps': num_inference_steps,
			    'output_type': output_type,
			}

			if image is not None:
					UpperCamelCase__				:					Dict					=					image

			if mask_image is not None:
					UpperCamelCase__				:					Optional[int]					=					mask_image

			if original_image is not None:
					UpperCamelCase__				:					Union[str, Any]					=					original_image

			# set all optional components to None
			for optional_component in pipe._optional_components:
					setattr(UpperCAmelCase_  ,					UpperCAmelCase_  ,					UpperCAmelCase_)

			UpperCamelCase__				:					int					=					pipe(**UpperCAmelCase_)[0]

			with tempfile.TemporaryDirectory() as tmpdir:
					pipe.save_pretrained(UpperCAmelCase_)
					UpperCamelCase__				:					Optional[Any]					=					self.pipeline_class.from_pretrained(UpperCAmelCase_)
					pipe_loaded.to(UpperCAmelCase_)
					pipe_loaded.set_progress_bar_config(disable=UpperCAmelCase_)

			pipe_loaded.unet.set_attn_processor(AttnAddedKVProcessor())  # For reproducibility tests

			for optional_component in pipe._optional_components:
					self.assertTrue(
					    getattr(UpperCAmelCase_  ,					UpperCAmelCase_) is None  ,					F'`{optional_component}` did not stay set to None after loading.'  ,					)

			UpperCamelCase__				:					Optional[int]					=					self.get_dummy_inputs(UpperCAmelCase_)

			UpperCamelCase__				:					Union[str, Any]					=					inputs['generator']
			UpperCamelCase__				:					List[Any]					=					inputs['num_inference_steps']
			UpperCamelCase__				:					Optional[int]					=					inputs['output_type']

			# inputs with prompt converted to embeddings
			UpperCamelCase__				:					Any					=					{
			    'prompt_embeds': prompt_embeds,
			    'negative_prompt_embeds': negative_prompt_embeds,
			    'generator': generator,
			    'num_inference_steps': num_inference_steps,
			    'output_type': output_type,
			}

			if image is not None:
					UpperCamelCase__				:					Tuple					=					image

			if mask_image is not None:
					UpperCamelCase__				:					Union[str, Any]					=					mask_image

			if original_image is not None:
					UpperCamelCase__				:					str					=					original_image

			UpperCamelCase__				:					Union[str, Any]					=					pipe_loaded(**UpperCAmelCase_)[0]

			UpperCamelCase__				:					Dict					=					np.abs(to_np(UpperCAmelCase_) - to_np(UpperCAmelCase_)).max()
			self.assertLess(UpperCAmelCase_  ,					1e-4)





	def      __UpperCamelCase							(      self   :						Optional[int]):
			UpperCamelCase__				:					Any					=					self.get_dummy_components()
			UpperCamelCase__				:					List[str]					=					self.pipeline_class(**UpperCAmelCase_)
			pipe.to(UpperCAmelCase_)
			pipe.set_progress_bar_config(disable=UpperCAmelCase_)

			UpperCamelCase__				:					Union[str, Any]					=					self.get_dummy_inputs(UpperCAmelCase_)
			UpperCamelCase__				:					Any					=					pipe(**UpperCAmelCase_)[0]

			with tempfile.TemporaryDirectory() as tmpdir:
					pipe.save_pretrained(UpperCAmelCase_)
					UpperCamelCase__				:					Optional[Any]					=					self.pipeline_class.from_pretrained(UpperCAmelCase_)
					pipe_loaded.to(UpperCAmelCase_)
					pipe_loaded.set_progress_bar_config(disable=UpperCAmelCase_)

			pipe_loaded.unet.set_attn_processor(AttnAddedKVProcessor())  # For reproducibility tests

			UpperCamelCase__				:					Any					=					self.get_dummy_inputs(UpperCAmelCase_)
			UpperCamelCase__				:					Tuple					=					pipe_loaded(**UpperCAmelCase_)[0]

			UpperCamelCase__				:					Optional[int]					=					np.abs(to_np(UpperCAmelCase_) - to_np(UpperCAmelCase_)).max()
			self.assertLess(UpperCAmelCase_  ,					1e-4)




