import gc
import random
import unittest

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

from diffusers import AutoencoderKL, PNDMScheduler, StableDiffusionInpaintPipeline, UNetaDConditionModel
from diffusers.utils import floats_tensor, load_image, load_numpy, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu, slow

from ..pipeline_params import TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS, TEXT_GUIDED_IMAGE_INPAINTING_PARAMS
from ..test_pipelines_common import PipelineKarrasSchedulerTesterMixin, PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()



class  UpperCamelCase				( _UpperCAmelCase   ,  _UpperCAmelCase   ,  _UpperCAmelCase   ,  unittest.TestCase ):

     '''simple docstring'''

     A_										=  StableDiffusionInpaintPipeline
     A_										=  TEXT_GUIDED_IMAGE_INPAINTING_PARAMS
     A_										=  TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS
     A_										=  frozenset(
         [] )  # TO-DO: update image_params once pipeline is refactored with VaeImageProcessor.preprocess
     A_										=  frozenset([] )



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






           """simple docstring"""
           torch.manual_seed(0						)
           _lowerCamelCase			    =				UNetaDConditionModel(
               block_out_channels=(32, 64)				,				layers_per_block=2				,				sample_size=32				,				in_channels=9				,				out_channels=4				,				down_block_types=('''DownBlock2D''', '''CrossAttnDownBlock2D''')				,				up_block_types=('''CrossAttnUpBlock2D''', '''UpBlock2D''')				,				cross_attention_dim=32				,				attention_head_dim=(2, 4)				,				use_linear_projection=__UpperCamelCase				,				)
           _lowerCamelCase			    =				PNDMScheduler(skip_prk_steps=__UpperCamelCase						)
           torch.manual_seed(0						)
           _lowerCamelCase			    =				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				,				sample_size=1_28				,				)
           torch.manual_seed(0						)
           _lowerCamelCase			    =				CLIPTextConfig(
               bos_token_id=0				,				eos_token_id=2				,				hidden_size=32				,				intermediate_size=37				,				layer_norm_eps=1E-0_5				,				num_attention_heads=4				,				num_hidden_layers=5				,				pad_token_id=1				,				vocab_size=10_00				,				hidden_act='''gelu'''				,				projection_dim=5_12				,				)
           _lowerCamelCase			    =				CLIPTextModel(__UpperCamelCase						)
           _lowerCamelCase			    =				CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''						)

           _lowerCamelCase			    =				{
               '''unet''': unet,
               '''scheduler''': scheduler,
               '''vae''': vae,
               '''text_encoder''': text_encoder,
               '''tokenizer''': tokenizer,
               '''safety_checker''': None,
               '''feature_extractor''': None,
           }
           return components



     def 		UpperCamelCase_						(					self				,				A_				,				A_=0						) ->				Union[str, Any]:






           """simple docstring"""
           _lowerCamelCase			    =				floats_tensor((1, 3, 32, 32)				,				rng=random.Random(__UpperCamelCase						)						).to(__UpperCamelCase						)
           _lowerCamelCase			    =				image.cpu().permute(0				,				2				,				3				,				1						)[0]
           _lowerCamelCase			    =				Image.fromarray(np.uinta(__UpperCamelCase						)						).convert('''RGB'''						).resize((64, 64)						)
           _lowerCamelCase			    =				Image.fromarray(np.uinta(image + 4						)						).convert('''RGB'''						).resize((64, 64)						)
           if str(__UpperCamelCase						).startswith('''mps'''						):
                 _lowerCamelCase			    =				torch.manual_seed(__UpperCamelCase						)
           else:
                 _lowerCamelCase			    =				torch.Generator(device=__UpperCamelCase						).manual_seed(__UpperCamelCase						)
           _lowerCamelCase			    =				{
               '''prompt''': '''A painting of a squirrel eating a burger''',
               '''image''': init_image,
               '''mask_image''': mask_image,
               '''generator''': generator,
               '''num_inference_steps''': 2,
               '''guidance_scale''': 6.0,
               '''output_type''': '''numpy''',
           }
           return inputs



     def 		UpperCamelCase_						(					self						) ->				int:






           """simple docstring"""
           _lowerCamelCase			    =				'''cpu'''  # ensure determinism for the device-dependent torch.Generator
           _lowerCamelCase			    =				self.get_dummy_components()
           _lowerCamelCase			    =				StableDiffusionInpaintPipeline(**__UpperCamelCase						)
           _lowerCamelCase			    =				sd_pipe.to(__UpperCamelCase						)
           sd_pipe.set_progress_bar_config(disable=__UpperCamelCase						)

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

           assert image.shape == (1, 64, 64, 3)
           _lowerCamelCase			    =				np.array([0.4727, 0.5735, 0.3941, 0.5446, 0.5926, 0.4394, 0.5062, 0.4654, 0.4476]						)

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





     def 		UpperCamelCase_						(					self						) ->				List[Any]:






           """simple docstring"""
           super().test_inference_batch_single_identical(expected_max_diff=3E-3						)





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

     '''simple docstring'''

     def 		UpperCamelCase_						(					self						) ->				Optional[Any]:






           """simple docstring"""
           super().tearDown()
           gc.collect()
           torch.cuda.empty_cache()



     def 		UpperCamelCase_						(					self						) ->				Dict:






           """simple docstring"""
           _lowerCamelCase			    =				load_image(
               '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
               '''/sd2-inpaint/init_image.png'''						)
           _lowerCamelCase			    =				load_image(
               '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd2-inpaint/mask.png'''						)
           _lowerCamelCase			    =				load_numpy(
               '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd2-inpaint'''
               '''/yellow_cat_sitting_on_a_park_bench.npy'''						)

           _lowerCamelCase			    =				'''stabilityai/stable-diffusion-2-inpainting'''
           _lowerCamelCase			    =				StableDiffusionInpaintPipeline.from_pretrained(__UpperCamelCase				,				safety_checker=__UpperCamelCase						)
           pipe.to(__UpperCamelCase						)
           pipe.set_progress_bar_config(disable=__UpperCamelCase						)
           pipe.enable_attention_slicing()

           _lowerCamelCase			    =				'''Face of a yellow cat, high resolution, sitting on a park bench'''

           _lowerCamelCase			    =				torch.manual_seed(0						)
           _lowerCamelCase			    =				pipe(
               prompt=__UpperCamelCase				,				image=__UpperCamelCase				,				mask_image=__UpperCamelCase				,				generator=__UpperCamelCase				,				output_type='''np'''				,				)
           _lowerCamelCase			    =				output.images[0]

           assert image.shape == (5_12, 5_12, 3)
           assert np.abs(expected_image - image						).max() < 9E-3



     def 		UpperCamelCase_						(					self						) ->				List[Any]:






           """simple docstring"""
           _lowerCamelCase			    =				load_image(
               '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
               '''/sd2-inpaint/init_image.png'''						)
           _lowerCamelCase			    =				load_image(
               '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd2-inpaint/mask.png'''						)
           _lowerCamelCase			    =				load_numpy(
               '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd2-inpaint'''
               '''/yellow_cat_sitting_on_a_park_bench_fp16.npy'''						)

           _lowerCamelCase			    =				'''stabilityai/stable-diffusion-2-inpainting'''
           _lowerCamelCase			    =				StableDiffusionInpaintPipeline.from_pretrained(
               __UpperCamelCase				,				torch_dtype=torch.floataa				,				safety_checker=__UpperCamelCase				,				)
           pipe.to(__UpperCamelCase						)
           pipe.set_progress_bar_config(disable=__UpperCamelCase						)
           pipe.enable_attention_slicing()

           _lowerCamelCase			    =				'''Face of a yellow cat, high resolution, sitting on a park bench'''

           _lowerCamelCase			    =				torch.manual_seed(0						)
           _lowerCamelCase			    =				pipe(
               prompt=__UpperCamelCase				,				image=__UpperCamelCase				,				mask_image=__UpperCamelCase				,				generator=__UpperCamelCase				,				output_type='''np'''				,				)
           _lowerCamelCase			    =				output.images[0]

           assert image.shape == (5_12, 5_12, 3)
           assert np.abs(expected_image - image						).max() < 5E-1





     def 		UpperCamelCase_						(					self						) ->				Any:






           """simple docstring"""
           torch.cuda.empty_cache()
           torch.cuda.reset_max_memory_allocated()
           torch.cuda.reset_peak_memory_stats()

           _lowerCamelCase			    =				load_image(
               '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
               '''/sd2-inpaint/init_image.png'''						)
           _lowerCamelCase			    =				load_image(
               '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd2-inpaint/mask.png'''						)

           _lowerCamelCase			    =				'''stabilityai/stable-diffusion-2-inpainting'''
           _lowerCamelCase			    =				PNDMScheduler.from_pretrained(__UpperCamelCase				,				subfolder='''scheduler'''						)
           _lowerCamelCase			    =				StableDiffusionInpaintPipeline.from_pretrained(
               __UpperCamelCase				,				safety_checker=__UpperCamelCase				,				scheduler=__UpperCamelCase				,				torch_dtype=torch.floataa				,				)
           pipe.to(__UpperCamelCase						)
           pipe.set_progress_bar_config(disable=__UpperCamelCase						)
           pipe.enable_attention_slicing(1						)
           pipe.enable_sequential_cpu_offload()

           _lowerCamelCase			    =				'''Face of a yellow cat, high resolution, sitting on a park bench'''

           _lowerCamelCase			    =				torch.manual_seed(0						)
           _lowerCamelCase			    =				pipe(
               prompt=__UpperCamelCase				,				image=__UpperCamelCase				,				mask_image=__UpperCamelCase				,				generator=__UpperCamelCase				,				num_inference_steps=2				,				output_type='''np'''				,				)

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


import argparse
import json
import os

import fairseq
import torch
from fairseq.data import Dictionary

from transformers import (
    UniSpeechConfig,
    UniSpeechForCTC,
    UniSpeechForPreTraining,
    WavaVecaFeatureExtractor,
    WavaVecaPhonemeCTCTokenizer,
    WavaVecaProcessor,
    logging,
)


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

snake_case__				= {
    'post_extract_proj': 'feature_projection.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.layer_norm': 'encoder.layer_norm',
    'w2v_model.layer_norm': 'feature_projection.layer_norm',
    'quantizer.weight_proj': 'quantizer.weight_proj',
    'quantizer.vars': 'quantizer.codevectors',
    'project_q': 'project_q',
    'final_proj': 'project_hid',
    'w2v_encoder.proj': 'ctc_proj',
    'mask_emb': 'masked_spec_embed',
}
snake_case__				= [
    'ctc_proj',
    'quantizer.weight_proj',
    'quantizer.codevectors',
    'project_q',
    'project_hid',
]
def       __magic_name__(  __UpperCAmelCase ,						__UpperCAmelCase ,						__UpperCAmelCase ,						__UpperCAmelCase ,						__UpperCAmelCase ,						__UpperCAmelCase		)						->    Optional[Any]:



      '''simple docstring'''

      for attribute in key.split('''.'''		):
            if is_finetuned:
                  if attribute in ["quantizer", "project_q", "project_hid"]:
                        # those layers are only relevant for pretraining and should be dropped
                        return

                  if attribute == "ctc_proj":
                        # we should rename `ctc_proj` to `lm_head` for fine-tuned phoneme models
                        _lowerCamelCase			    =				'''lm_head'''

            _lowerCamelCase			    =				getattr(__UpperCAmelCase ,						__UpperCAmelCase		)

      if weight_type is not None:
            _lowerCamelCase			    =				getattr(__UpperCAmelCase ,						__UpperCAmelCase		).shape
      else:
            _lowerCamelCase			    =				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":
            _lowerCamelCase			    =				value
      elif weight_type == "weight_g":
            _lowerCamelCase			    =				value
      elif weight_type == "weight_v":
            _lowerCamelCase			    =				value
      elif weight_type == "bias":
            _lowerCamelCase			    =				value
      else:
            _lowerCamelCase			    =				value

      logger.info(F'{key + "." + weight_type if weight_type is not None else ""} was initialized from {full_name}.'		)
def       __magic_name__(  __UpperCAmelCase ,						__UpperCAmelCase ,						__UpperCAmelCase		)						->    List[str]:



      '''simple docstring'''

      _lowerCamelCase			    =				[]
      _lowerCamelCase			    =				fairseq_model.state_dict()

      _lowerCamelCase			    =				hf_model.unispeech.feature_extractor

      for name, value in fairseq_dict.items():
            _lowerCamelCase			    =				False
            if "conv_layers" in name:
                  load_conv_layer(
                      __UpperCAmelCase ,						__UpperCAmelCase ,						__UpperCAmelCase ,						__UpperCAmelCase ,						hf_model.config.feat_extract_norm == '''group''' ,						)
                  _lowerCamelCase			    =				True
            else:
                  for key, mapped_key in MAPPING.items():
                        _lowerCamelCase			    =				'''unispeech.''' + mapped_key if mapped_key not in TOP_LEVEL_KEYS else mapped_key
                        if key in name or key.split('''w2v_model.'''		)[-1] == name.split('''.'''		)[0]:
                              _lowerCamelCase			    =				True
                              if "*" in mapped_key:
                                    _lowerCamelCase			    =				name.split(__UpperCAmelCase		)[0].split('''.'''		)[-2]
                                    _lowerCamelCase			    =				mapped_key.replace('''*''' ,						__UpperCAmelCase		)
                              if "weight_g" in name:
                                    _lowerCamelCase			    =				'''weight_g'''
                              elif "weight_v" in name:
                                    _lowerCamelCase			    =				'''weight_v'''
                              elif "bias" in name:
                                    _lowerCamelCase			    =				'''bias'''
                              elif "weight" in name:
                                    # TODO: don't match quantizer.weight_proj
                                    _lowerCamelCase			    =				'''weight'''
                              else:
                                    _lowerCamelCase			    =				None
                              set_recursively(__UpperCAmelCase ,						__UpperCAmelCase ,						__UpperCAmelCase ,						__UpperCAmelCase ,						__UpperCAmelCase ,						__UpperCAmelCase		)
                        continue
            if not is_used:
                  unused_weights.append(__UpperCAmelCase		)

      logger.warning(F'Unused weights: {unused_weights}'		)
def       __magic_name__(  __UpperCAmelCase ,						__UpperCAmelCase ,						__UpperCAmelCase ,						__UpperCAmelCase ,						__UpperCAmelCase		)						->    int:



      '''simple docstring'''

      _lowerCamelCase			    =				full_name.split('''conv_layers.'''		)[-1]
      _lowerCamelCase			    =				name.split('''.'''		)
      _lowerCamelCase			    =				int(items[0]		)
      _lowerCamelCase			    =				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.'
                  )
                  _lowerCamelCase			    =				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.'
                  )
                  _lowerCamelCase			    =				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."
                  )
                  _lowerCamelCase			    =				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.'
                  )
                  _lowerCamelCase			    =				value
                  logger.info(F'Feat extract layer norm weight of layer {layer_id} was initialized from {full_name}.'		)
      else:
            unused_weights.append(__UpperCAmelCase		)
@torch.no_grad()
def       __magic_name__(  __UpperCAmelCase ,						__UpperCAmelCase ,						__UpperCAmelCase=None ,						__UpperCAmelCase=None ,						__UpperCAmelCase=True		)						->    Union[str, Any]:



      '''simple docstring'''

      if config_path is not None:
            _lowerCamelCase			    =				UniSpeechConfig.from_pretrained(__UpperCAmelCase		)
      else:
            _lowerCamelCase			    =				UniSpeechConfig()

      if is_finetuned:
            if dict_path:
                  _lowerCamelCase			    =				Dictionary.load_from_json(__UpperCAmelCase		)

                  # important change bos & pad token id since CTC symbol is <pad> and
                  # not <s> as in fairseq
                  _lowerCamelCase			    =				target_dict.pad_index
                  _lowerCamelCase			    =				target_dict.bos_index
                  _lowerCamelCase			    =				target_dict.eos_index
                  _lowerCamelCase			    =				len(target_dict.symbols		)
                  _lowerCamelCase			    =				os.path.join(__UpperCAmelCase ,						'''vocab.json'''		)
                  if not os.path.isdir(__UpperCAmelCase		):
                        logger.error('''--pytorch_dump_folder_path ({}) should be a directory'''.format(__UpperCAmelCase		)		)
                        return
                  os.makedirs(__UpperCAmelCase ,						exist_ok=__UpperCAmelCase		)
                  _lowerCamelCase			    =				target_dict.indices

                  # fairseq has the <pad> and <s> switched
                  _lowerCamelCase			    =				42
                  _lowerCamelCase			    =				43
                  with open(__UpperCAmelCase ,						'''w''' ,						encoding='''utf-8'''		) as vocab_handle:
                        json.dump(__UpperCAmelCase ,						__UpperCAmelCase		)
                  _lowerCamelCase			    =				WavaVecaPhonemeCTCTokenizer(
                      __UpperCAmelCase ,						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=__UpperCAmelCase ,						)
                  _lowerCamelCase			    =				True if config.feat_extract_norm == '''layer''' else False
                  _lowerCamelCase			    =				WavaVecaFeatureExtractor(
                      feature_size=1 ,						sampling_rate=1_6000 ,						padding_value=0 ,						do_normalize=__UpperCAmelCase ,						return_attention_mask=__UpperCAmelCase ,						)
                  _lowerCamelCase			    =				WavaVecaProcessor(feature_extractor=__UpperCAmelCase ,						tokenizer=__UpperCAmelCase		)
                  processor.save_pretrained(__UpperCAmelCase		)

            _lowerCamelCase			    =				UniSpeechForCTC(__UpperCAmelCase		)
      else:
            _lowerCamelCase			    =				UniSpeechForPreTraining(__UpperCAmelCase		)

      if is_finetuned:
            _lowerCamelCase    ,					_lowerCamelCase    ,					_lowerCamelCase			    =				fairseq.checkpoint_utils.load_model_ensemble_and_task(
                [checkpoint_path] ,						arg_overrides={'''data''': '''/'''.join(dict_path.split('''/'''		)[:-1]		), '''w2v_path''': checkpoint_path}		)
      else:
            _lowerCamelCase    ,					_lowerCamelCase    ,					_lowerCamelCase			    =				fairseq.checkpoint_utils.load_model_ensemble_and_task([checkpoint_path]		)

      _lowerCamelCase			    =				model[0].eval()

      recursively_load_weights(__UpperCAmelCase ,						__UpperCAmelCase ,						__UpperCAmelCase		)

      hf_unispeech.save_pretrained(__UpperCAmelCase		)


if __name__ == "__main__":
     snake_case__				= 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(
         '--not_finetuned', action='store_true', help='Whether the model to convert is a fine-tuned model or not'
     )
     snake_case__				= parser.parse_args()
     convert_unispeech_checkpoint(
         args.checkpoint_path, args.pytorch_dump_folder_path, args.config_path, args.dict_path, not args.not_finetuned
     )