'''simple docstring'''




import gc
import random
import unittest

import numpy as np
import torch
from PIL import Image

from diffusers import (
    DDIMScheduler,
    KandinskyVaaControlnetImgaImgPipeline,
    KandinskyVaaPriorEmbaEmbPipeline,
    UNetaDConditionModel,
    VQModel,
)
from diffusers.utils import floats_tensor, load_image, load_numpy, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..test_pipelines_common import PipelineTesterMixin, assert_mean_pixel_difference


enable_full_determinism()




class    a__				(       UpperCAmelCase__							,       unittest.TestCase			):
     lowerCamelCase  :					int							=KandinskyVaaControlnetImgaImgPipeline
     lowerCamelCase  :					List[str]							=["image_embeds", "negative_image_embeds", "image", "hint"]
     lowerCamelCase  :					Dict							=["image_embeds", "negative_image_embeds", "image", "hint"]
     lowerCamelCase  :					Any							=[
         "generator",
         "height",
         "width",
         "strength",
         "guidance_scale",
         "num_inference_steps",
         "return_dict",
         "guidance_scale",
         "num_images_per_prompt",
         "output_type",
         "return_dict",
     ]
     lowerCamelCase  :					Tuple							=False






     @property
     def        SCREAMING_SNAKE_CASE__       (			self		:		str   ):





         """simple docstring"""



         return 32






     @property
     def        SCREAMING_SNAKE_CASE__       (			self		:		int   ):





         """simple docstring"""



         return 32






     @property
     def        SCREAMING_SNAKE_CASE__       (			self		:		Any   ):





         """simple docstring"""



         return self.time_input_dim






     @property
     def        SCREAMING_SNAKE_CASE__       (			self		:		Optional[int]   ):





         """simple docstring"""



         return self.time_input_dim * 4






     @property
     def        SCREAMING_SNAKE_CASE__       (			self		:		Union[str, Any]   ):





         """simple docstring"""



         return 1_00






     @property
     def        SCREAMING_SNAKE_CASE__       (			self		:		Dict   ):





         """simple docstring"""



         torch.manual_seed(0   )

         __lowerCamelCase          =  {
             '''in_channels''': 8,
             # Out channels is double in channels because predicts mean and variance
             '''out_channels''': 8,
             '''addition_embed_type''': '''image_hint''',
             '''down_block_types''': ('''ResnetDownsampleBlock2D''', '''SimpleCrossAttnDownBlock2D'''),
             '''up_block_types''': ('''SimpleCrossAttnUpBlock2D''', '''ResnetUpsampleBlock2D'''),
             '''mid_block_type''': '''UNetMidBlock2DSimpleCrossAttn''',
             '''block_out_channels''': (self.block_out_channels_a, self.block_out_channels_a * 2),
             '''layers_per_block''': 1,
             '''encoder_hid_dim''': self.text_embedder_hidden_size,
             '''encoder_hid_dim_type''': '''image_proj''',
             '''cross_attention_dim''': self.cross_attention_dim,
             '''attention_head_dim''': 4,
             '''resnet_time_scale_shift''': '''scale_shift''',
             '''class_embed_type''': None,
         }

         __lowerCamelCase          =  UNetaDConditionModel(**a   )
         return model






     @property
     def        SCREAMING_SNAKE_CASE__       (			self		:		Any   ):





         """simple docstring"""



         return {
             "block_out_channels": [32, 32, 64, 64],
             "down_block_types": [
                 "DownEncoderBlock2D",
                 "DownEncoderBlock2D",
                 "DownEncoderBlock2D",
                 "AttnDownEncoderBlock2D",
             ],
             "in_channels": 3,
             "latent_channels": 4,
             "layers_per_block": 1,
             "norm_num_groups": 8,
             "norm_type": "spatial",
             "num_vq_embeddings": 12,
             "out_channels": 3,
             "up_block_types": ["AttnUpDecoderBlock2D", "UpDecoderBlock2D", "UpDecoderBlock2D", "UpDecoderBlock2D"],
             "vq_embed_dim": 4,
         }






     @property
     def        SCREAMING_SNAKE_CASE__       (			self		:		Optional[Any]   ):





         """simple docstring"""



         torch.manual_seed(0   )
         __lowerCamelCase          =  VQModel(**self.dummy_movq_kwargs   )
         return model






     def        SCREAMING_SNAKE_CASE__       (			self		:		Optional[Any]   ):





         """simple docstring"""



         __lowerCamelCase          =  self.dummy_unet
         __lowerCamelCase          =  self.dummy_movq

         __lowerCamelCase          =  {
             '''num_train_timesteps''': 10_00,
             '''beta_schedule''': '''linear''',
             '''beta_start''': 0.0_00_85,
             '''beta_end''': 0.0_12,
             '''clip_sample''': False,
             '''set_alpha_to_one''': False,
             '''steps_offset''': 0,
             '''prediction_type''': '''epsilon''',
             '''thresholding''': False,
         }

         __lowerCamelCase          =  DDIMScheduler(**a   )

         __lowerCamelCase          =  {
             '''unet''': unet,
             '''scheduler''': scheduler,
             '''movq''': movq,
         }

         return components






     def        SCREAMING_SNAKE_CASE__       (			self		:		List[str]  ,    a		:		str  ,    a		:		Union[str, Any]=0   ):





         """simple docstring"""



         __lowerCamelCase          =  floats_tensor((1, self.text_embedder_hidden_size)  ,    rng=random.Random(a   )   ).to(a   )
         __lowerCamelCase          =  floats_tensor((1, self.text_embedder_hidden_size)  ,    rng=random.Random(seed + 1   )   ).to(
             a   )
         # create init_image
         __lowerCamelCase          =  floats_tensor((1, 3, 64, 64)  ,    rng=random.Random(a   )   ).to(a   )
         __lowerCamelCase          =  image.cpu().permute(0  ,    2  ,    3  ,    1   )[0]
         __lowerCamelCase          =  Image.fromarray(np.uinta(a   )   ).convert('''RGB'''   ).resize((2_56, 2_56)   )
         # create hint
         __lowerCamelCase          =  floats_tensor((1, 3, 64, 64)  ,    rng=random.Random(a   )   ).to(a   )

         if str(a   ).startswith('''mps'''   ):
             __lowerCamelCase          =  torch.manual_seed(a   )
         else:
             __lowerCamelCase          =  torch.Generator(device=a   ).manual_seed(a   )
         __lowerCamelCase          =  {
             '''image''': init_image,
             '''image_embeds''': image_embeds,
             '''negative_image_embeds''': negative_image_embeds,
             '''hint''': hint,
             '''generator''': generator,
             '''height''': 64,
             '''width''': 64,
             '''num_inference_steps''': 10,
             '''guidance_scale''': 7.0,
             '''strength''': 0.2,
             '''output_type''': '''np''',
         }
         return inputs






     def        SCREAMING_SNAKE_CASE__       (			self		:		str   ):





         """simple docstring"""



         __lowerCamelCase          =  '''cpu'''

         __lowerCamelCase          =  self.get_dummy_components()

         __lowerCamelCase          =  self.pipeline_class(**a   )
         __lowerCamelCase          =  pipe.to(a   )

         pipe.set_progress_bar_config(disable=a   )

         __lowerCamelCase          =  pipe(**self.get_dummy_inputs(a   )   )
         __lowerCamelCase          =  output.images

         __lowerCamelCase          =  pipe(
             **self.get_dummy_inputs(a   )  ,    return_dict=a  ,    )[0]

         __lowerCamelCase          =  image[0, -3:, -3:, -1]
         __lowerCamelCase          =  image_from_tuple[0, -3:, -3:, -1]

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

         __lowerCamelCase          =  np.array(
             [0.54_98_50_34, 0.55_50_93_65, 0.52_56_15_04, 0.5_57_04_94, 0.5_59_38_18, 0.5_26_39_79, 0.50_28_56_43, 0.5_06_98_46, 0.51_19_67_36]   )
         assert (
             np.abs(image_slice.flatten() - expected_slice   ).max() < 1e-2
         ), f""" expected_slice {expected_slice}, but got {image_slice.flatten()}"""
         assert (
             np.abs(image_from_tuple_slice.flatten() - expected_slice   ).max() < 1e-2
         ), f""" expected_slice {expected_slice}, but got {image_from_tuple_slice.flatten()}"""




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






     def        SCREAMING_SNAKE_CASE__       (			self		:		Optional[int]   ):





         """simple docstring"""



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






     def        SCREAMING_SNAKE_CASE__       (			self		:		Optional[int]   ):





         """simple docstring"""



         __lowerCamelCase          =  load_numpy(
             '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
             '''/kandinskyv22/kandinskyv22_controlnet_img2img_robotcat_fp16.npy'''   )

         __lowerCamelCase          =  load_image(
             '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main''' '''/kandinsky/cat.png'''   )
         __lowerCamelCase          =  init_image.resize((5_12, 5_12)   )

         __lowerCamelCase          =  load_image(
             '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
             '''/kandinskyv22/hint_image_cat.png'''   )
         __lowerCamelCase          =  torch.from_numpy(np.array(a   )   ).float() / 2_55.0
         __lowerCamelCase          =  hint.permute(2  ,    0  ,    1   ).unsqueeze(0   )

         __lowerCamelCase          =  '''A robot, 4k photo'''

         __lowerCamelCase          =  KandinskyVaaPriorEmbaEmbPipeline.from_pretrained(
             '''kandinsky-community/kandinsky-2-2-prior'''  ,    torch_dtype=torch.floataa   )
         pipe_prior.to(a   )

         __lowerCamelCase          =  KandinskyVaaControlnetImgaImgPipeline.from_pretrained(
             '''kandinsky-community/kandinsky-2-2-controlnet-depth'''  ,    torch_dtype=torch.floataa   )
         __lowerCamelCase          =  pipeline.to(a   )

         pipeline.set_progress_bar_config(disable=a   )

         __lowerCamelCase          =  torch.Generator(device='''cpu'''   ).manual_seed(0   )

         __lowerCamelCase     ,  __lowerCamelCase          =  pipe_prior(
             a  ,    image=a  ,    strength=0.85  ,    generator=a  ,    negative_prompt=''''''  ,    ).to_tuple()

         __lowerCamelCase          =  pipeline(
             image=a  ,    image_embeds=a  ,    negative_image_embeds=a  ,    hint=a  ,    generator=a  ,    num_inference_steps=1_00  ,    height=5_12  ,    width=5_12  ,    strength=0.5  ,    output_type='''np'''  ,    )

         __lowerCamelCase          =  output.images[0]

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

         assert_mean_pixel_difference(a  ,    a   )

'''simple docstring'''




import argparse

import fairseq
import torch

from transformers import UniSpeechSatConfig, UniSpeechSatForCTC, UniSpeechSatForPreTraining, logging


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

__UpperCAmelCase												={
    "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",
    "encoder.layer_norm_for_extract": "layer_norm_for_extract",
    "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": "lm_head",
    "label_embs_concat": "label_embeddings_concat",
    "mask_emb": "masked_spec_embed",
    "spk_proj": "speaker_proj",
}
__UpperCAmelCase												=[
    "lm_head",
    "quantizer.weight_proj",
    "quantizer.codevectors",
    "project_q",
    "project_hid",
    "label_embeddings_concat",
    "speaker_proj",
    "layer_norm_for_extract",
]






def __lowerCAmelCase			(							UpperCamelCase__    ,  UpperCamelCase__    ,  UpperCamelCase__    ,  UpperCamelCase__    ,  UpperCamelCase__    )     ->   Any:
    for attribute in key.split('''.'''    ):
        __lowerCamelCase          =  getattr(UpperCamelCase__    ,  UpperCamelCase__    )

    if weight_type is not None:
        __lowerCamelCase          =  getattr(UpperCamelCase__    ,  UpperCamelCase__    ).shape
    else:
        __lowerCamelCase          =  hf_pointer.shape

    if hf_shape != value.shape:
        raise ValueError(
            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 __lowerCAmelCase			(							UpperCamelCase__    ,  UpperCamelCase__    )     ->   List[Any]:
    __lowerCamelCase          =  []
    __lowerCamelCase          =  fairseq_model.state_dict()

    __lowerCamelCase          =  hf_model.unispeech_sat.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_sat.''' + 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]:
                    if "layer_norm_for_extract" in name and (".".join(name.split('''.'''    )[:-1]    ) != key):
                        # special case since naming is very similar
                        continue
                    __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__    )
                continue
        if not is_used:
            unused_weights.append(UpperCamelCase__    )

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






def __lowerCAmelCase			(							UpperCamelCase__    ,  UpperCamelCase__    ,  UpperCamelCase__    ,  UpperCamelCase__    ,  UpperCamelCase__    )     ->   Any:
    __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:
            if value.shape != feature_extractor.conv_layers[layer_id].conv.bias.data.shape:
                raise ValueError(
                    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:
            if value.shape != feature_extractor.conv_layers[layer_id].conv.weight.data.shape:
                raise ValueError(
                    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:
            if value.shape != feature_extractor.conv_layers[layer_id].layer_norm.bias.data.shape:
                raise ValueError(
                    f"""{full_name} has size {value.shape}, but"""
                    f""" {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:
            if value.shape != feature_extractor.conv_layers[layer_id].layer_norm.weight.data.shape:
                raise ValueError(
                    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 __lowerCAmelCase			(							UpperCamelCase__    ,  UpperCamelCase__    ,  UpperCamelCase__=None    ,  UpperCamelCase__=None    ,  UpperCamelCase__=True    )     ->   Optional[Any]:
    if config_path is not None:
        __lowerCamelCase          =  UniSpeechSatConfig.from_pretrained(UpperCamelCase__    )
    else:
        __lowerCamelCase          =  UniSpeechSatConfig()

    __lowerCamelCase          =  ''''''

    if is_finetuned:
        __lowerCamelCase          =  UniSpeechSatForCTC(UpperCamelCase__    )
    else:
        __lowerCamelCase          =  UniSpeechSatForPreTraining(UpperCamelCase__    )

    __lowerCamelCase     ,  __lowerCamelCase     ,  __lowerCamelCase          =  fairseq.checkpoint_utils.load_model_ensemble_and_task(
        [checkpoint_path]    ,  arg_overrides={'''data''': '''/'''.join(dict_path.split('''/'''    )[:-1]    )}    )
    __lowerCamelCase          =  model[0].eval()

    recursively_load_weights(UpperCamelCase__    ,  UpperCamelCase__    )

    hf_wavavec.save_pretrained(UpperCamelCase__    )


if __name__ == "__main__":
  __UpperCAmelCase												=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"
  )
  __UpperCAmelCase												=parser.parse_args()
  convert_unispeech_sat_checkpoint(
      args.checkpoint_path, args.pytorch_dump_folder_path, args.config_path, args.dict_path, not args.not_finetuned
  )
