


'''simple docstring'''






import argparse

import torch

from transformers import (
    EncodecConfig,
    EncodecFeatureExtractor,
    EncodecModel,
    logging,
)


# checkpoints downloaded from:
# https://dl.fbaipublicfiles.com/encodec/v0/encodec_24khz-d7cc33bc.th
# https://huggingface.co/facebook/musicgen-small/resolve/main/compression_state_dict.bin
# https://dl.fbaipublicfiles.com/encodec/v0/encodec_48khz-7e698e3e.th


logging.set_verbosity_info()
__A :  int    							=				logging.get_logger("transformers.models.encodec")

__A :  str    							=				{
    'quantizer.vq.layers.*._codebook.inited': 'quantizer.layers.*.codebook.inited',
    'quantizer.vq.layers.*._codebook.cluster_size': 'quantizer.layers.*.codebook.cluster_size',
    'quantizer.vq.layers.*._codebook.embed': 'quantizer.layers.*.codebook.embed',
    'quantizer.vq.layers.*._codebook.embed_avg': 'quantizer.layers.*.codebook.embed_avg',
}
__A :  Union[str, Any]    							=				{
    'encoder.model.0.conv.conv': 'encoder.layers.0.conv',
    'encoder.model.1.block.1.conv.conv': 'encoder.layers.1.block.1.conv',
    'encoder.model.1.block.3.conv.conv': 'encoder.layers.1.block.3.conv',
    'encoder.model.1.shortcut.conv.conv': 'encoder.layers.1.shortcut.conv',
    'encoder.model.3.conv.conv': 'encoder.layers.3.conv',
    'encoder.model.4.block.1.conv.conv': 'encoder.layers.4.block.1.conv',
    'encoder.model.4.block.3.conv.conv': 'encoder.layers.4.block.3.conv',
    'encoder.model.4.shortcut.conv.conv': 'encoder.layers.4.shortcut.conv',
    'encoder.model.6.conv.conv': 'encoder.layers.6.conv',
    'encoder.model.7.block.1.conv.conv': 'encoder.layers.7.block.1.conv',
    'encoder.model.7.block.3.conv.conv': 'encoder.layers.7.block.3.conv',
    'encoder.model.7.shortcut.conv.conv': 'encoder.layers.7.shortcut.conv',
    'encoder.model.9.conv.conv': 'encoder.layers.9.conv',
    'encoder.model.10.block.1.conv.conv': 'encoder.layers.10.block.1.conv',
    'encoder.model.10.block.3.conv.conv': 'encoder.layers.10.block.3.conv',
    'encoder.model.10.shortcut.conv.conv': 'encoder.layers.10.shortcut.conv',
    'encoder.model.12.conv.conv': 'encoder.layers.12.conv',
    'encoder.model.13.lstm': 'encoder.layers.13.lstm',
    'encoder.model.15.conv.conv': 'encoder.layers.15.conv',
}
__A :  Optional[int]    							=				{
    'encoder.model.0.conv.norm': 'encoder.layers.0.norm',
    'encoder.model.1.block.1.conv.norm': 'encoder.layers.1.block.1.norm',
    'encoder.model.1.block.3.conv.norm': 'encoder.layers.1.block.3.norm',
    'encoder.model.1.shortcut.conv.norm': 'encoder.layers.1.shortcut.norm',
    'encoder.model.3.conv.norm': 'encoder.layers.3.norm',
    'encoder.model.4.block.1.conv.norm': 'encoder.layers.4.block.1.norm',
    'encoder.model.4.block.3.conv.norm': 'encoder.layers.4.block.3.norm',
    'encoder.model.4.shortcut.conv.norm': 'encoder.layers.4.shortcut.norm',
    'encoder.model.6.conv.norm': 'encoder.layers.6.norm',
    'encoder.model.7.block.1.conv.norm': 'encoder.layers.7.block.1.norm',
    'encoder.model.7.block.3.conv.norm': 'encoder.layers.7.block.3.norm',
    'encoder.model.7.shortcut.conv.norm': 'encoder.layers.7.shortcut.norm',
    'encoder.model.9.conv.norm': 'encoder.layers.9.norm',
    'encoder.model.10.block.1.conv.norm': 'encoder.layers.10.block.1.norm',
    'encoder.model.10.block.3.conv.norm': 'encoder.layers.10.block.3.norm',
    'encoder.model.10.shortcut.conv.norm': 'encoder.layers.10.shortcut.norm',
    'encoder.model.12.conv.norm': 'encoder.layers.12.norm',
    'encoder.model.15.conv.norm': 'encoder.layers.15.norm',
}
__A :  int    							=				{
    'decoder.model.0.conv.conv': 'decoder.layers.0.conv',
    'decoder.model.1.lstm': 'decoder.layers.1.lstm',
    'decoder.model.3.convtr.convtr': 'decoder.layers.3.conv',
    'decoder.model.4.block.1.conv.conv': 'decoder.layers.4.block.1.conv',
    'decoder.model.4.block.3.conv.conv': 'decoder.layers.4.block.3.conv',
    'decoder.model.4.shortcut.conv.conv': 'decoder.layers.4.shortcut.conv',
    'decoder.model.6.convtr.convtr': 'decoder.layers.6.conv',
    'decoder.model.7.block.1.conv.conv': 'decoder.layers.7.block.1.conv',
    'decoder.model.7.block.3.conv.conv': 'decoder.layers.7.block.3.conv',
    'decoder.model.7.shortcut.conv.conv': 'decoder.layers.7.shortcut.conv',
    'decoder.model.9.convtr.convtr': 'decoder.layers.9.conv',
    'decoder.model.10.block.1.conv.conv': 'decoder.layers.10.block.1.conv',
    'decoder.model.10.block.3.conv.conv': 'decoder.layers.10.block.3.conv',
    'decoder.model.10.shortcut.conv.conv': 'decoder.layers.10.shortcut.conv',
    'decoder.model.12.convtr.convtr': 'decoder.layers.12.conv',
    'decoder.model.13.block.1.conv.conv': 'decoder.layers.13.block.1.conv',
    'decoder.model.13.block.3.conv.conv': 'decoder.layers.13.block.3.conv',
    'decoder.model.13.shortcut.conv.conv': 'decoder.layers.13.shortcut.conv',
    'decoder.model.15.conv.conv': 'decoder.layers.15.conv',
}
__A :  str    							=				{
    'decoder.model.0.conv.norm': 'decoder.layers.0.norm',
    'decoder.model.3.convtr.norm': 'decoder.layers.3.norm',
    'decoder.model.4.block.1.conv.norm': 'decoder.layers.4.block.1.norm',
    'decoder.model.4.block.3.conv.norm': 'decoder.layers.4.block.3.norm',
    'decoder.model.4.shortcut.conv.norm': 'decoder.layers.4.shortcut.norm',
    'decoder.model.6.convtr.norm': 'decoder.layers.6.norm',
    'decoder.model.7.block.1.conv.norm': 'decoder.layers.7.block.1.norm',
    'decoder.model.7.block.3.conv.norm': 'decoder.layers.7.block.3.norm',
    'decoder.model.7.shortcut.conv.norm': 'decoder.layers.7.shortcut.norm',
    'decoder.model.9.convtr.norm': 'decoder.layers.9.norm',
    'decoder.model.10.block.1.conv.norm': 'decoder.layers.10.block.1.norm',
    'decoder.model.10.block.3.conv.norm': 'decoder.layers.10.block.3.norm',
    'decoder.model.10.shortcut.conv.norm': 'decoder.layers.10.shortcut.norm',
    'decoder.model.12.convtr.norm': 'decoder.layers.12.norm',
    'decoder.model.13.block.1.conv.norm': 'decoder.layers.13.block.1.norm',
    'decoder.model.13.block.3.conv.norm': 'decoder.layers.13.block.3.norm',
    'decoder.model.13.shortcut.conv.norm': 'decoder.layers.13.shortcut.norm',
    'decoder.model.15.conv.norm': 'decoder.layers.15.norm',
}
__A :  Optional[Any]    							=				{
    **MAPPING_QUANTIZER,
    **MAPPING_ENCODER,
    **MAPPING_DECODER,
}
__A :  Tuple    							=				{
    **MAPPING_QUANTIZER,
    **MAPPING_ENCODER,
    **MAPPING_ENCODER_48K,
    **MAPPING_DECODER,
    **MAPPING_DECODER_48K,
}
__A :  Optional[Any]    							=				[]
__A :  str    							=				[]







def     UpperCamelCase_      (   A__ :       Union[str, Any]							,   A__ :       Tuple							,   A__ :       Any							,   A__ :       str							,   A__ :       str ):
		'''simple docstring'''





		for attribute in key.split(""".""" ):
				lowerCAmelCase_      :				List[Any]      =       getattr(lowercase__							,   lowercase__ )

		if weight_type is not None:
				lowerCAmelCase_      :				List[str]      =       getattr(lowercase__							,   lowercase__ ).shape
		else:
				lowerCAmelCase_      :				str      =       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_      :				List[Any]      =       value
		elif weight_type == "weight_g":
				lowerCAmelCase_      :				Optional[int]      =       value
		elif weight_type == "weight_v":
				lowerCAmelCase_      :				Tuple      =       value
		elif weight_type == "bias":
				lowerCAmelCase_      :				List[str]      =       value
		elif weight_type == "running_mean":
				lowerCAmelCase_      :				int      =       value
		elif weight_type == "running_var":
				lowerCAmelCase_      :				str      =       value
		elif weight_type == "num_batches_tracked":
				lowerCAmelCase_      :				List[str]      =       value
		elif weight_type == "weight_ih_l0":
				lowerCAmelCase_      :				List[str]      =       value
		elif weight_type == "weight_hh_l0":
				lowerCAmelCase_      :				Optional[int]      =       value
		elif weight_type == "bias_ih_l0":
				lowerCAmelCase_      :				Optional[int]      =       value
		elif weight_type == "bias_hh_l0":
				lowerCAmelCase_      :				Any      =       value
		elif weight_type == "weight_ih_l1":
				lowerCAmelCase_      :				Optional[Any]      =       value
		elif weight_type == "weight_hh_l1":
				lowerCAmelCase_      :				Dict      =       value
		elif weight_type == "bias_ih_l1":
				lowerCAmelCase_      :				str      =       value
		elif weight_type == "bias_hh_l1":
				lowerCAmelCase_      :				List[Any]      =       value
		else:
				lowerCAmelCase_      :				int      =       value

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







def     UpperCamelCase_      (   A__ :       str							,   A__ :       List[Any] ):
		'''simple docstring'''





		for key in ignore_keys:
				if key.endswith(""".*""" ):
						if name.startswith(key[:-1] ):
								return True
				elif ".*." in key:
						lowerCAmelCase_,							lowerCAmelCase_      :				Optional[Any]      =       key.split(""".*.""" )
						if prefix in name and suffix in name:
								return True
				elif key in name:
						return True
		return False







def     UpperCamelCase_      (   A__ :       List[str]							,   A__ :       str							,   A__ :       Optional[int] ):
		'''simple docstring'''





		lowerCAmelCase_      :				Optional[Any]      =       []

		if model_name == "encodec_24khz" or "encodec_32khz":
				lowerCAmelCase_      :				Optional[Any]      =       MAPPING_24K
		elif model_name == "encodec_48khz":
				lowerCAmelCase_      :				List[str]      =       MAPPING_48K
		else:
				raise ValueError(f'Unsupported model: {model_name}' )

		for name, value in orig_dict.items():
				if should_ignore(lowercase__							,   lowercase__ ):
						logger.info(f'{name} was ignored' )
						continue

				lowerCAmelCase_      :				Any      =       False
				for key, mapped_key in MAPPING.items():
						if "*" in key:
								lowerCAmelCase_,							lowerCAmelCase_      :				Union[str, Any]      =       key.split(""".*.""" )
								if prefix in name and suffix in name:
										lowerCAmelCase_      :				Any      =       suffix

						if key in name:
								# HACK otherwise .embed gets initialized with .embed_avg too
								if key.endswith("""embed""" ) and name.endswith("""embed_avg""" ):
										continue

								lowerCAmelCase_      :				Dict      =       True
								if "*" in mapped_key:
										lowerCAmelCase_      :				List[Any]      =       name.split(lowercase__ )[0].split(""".""" )[-2]
										lowerCAmelCase_      :				str      =       mapped_key.replace("""*"""							,   lowercase__ )
								if "weight_g" in name:
										lowerCAmelCase_      :				Dict      =       """weight_g"""
								elif "weight_v" in name:
										lowerCAmelCase_      :				List[str]      =       """weight_v"""
								elif "weight_ih_l0" in name:
										lowerCAmelCase_      :				Dict      =       """weight_ih_l0"""
								elif "weight_hh_l0" in name:
										lowerCAmelCase_      :				Tuple      =       """weight_hh_l0"""
								elif "bias_ih_l0" in name:
										lowerCAmelCase_      :				List[str]      =       """bias_ih_l0"""
								elif "bias_hh_l0" in name:
										lowerCAmelCase_      :				str      =       """bias_hh_l0"""
								elif "weight_ih_l1" in name:
										lowerCAmelCase_      :				str      =       """weight_ih_l1"""
								elif "weight_hh_l1" in name:
										lowerCAmelCase_      :				Union[str, Any]      =       """weight_hh_l1"""
								elif "bias_ih_l1" in name:
										lowerCAmelCase_      :				str      =       """bias_ih_l1"""
								elif "bias_hh_l1" in name:
										lowerCAmelCase_      :				Optional[int]      =       """bias_hh_l1"""
								elif "bias" in name:
										lowerCAmelCase_      :				Tuple      =       """bias"""
								elif "weight" in name:
										lowerCAmelCase_      :				Union[str, Any]      =       """weight"""
								elif "running_mean" in name:
										lowerCAmelCase_      :				int      =       """running_mean"""
								elif "running_var" in name:
										lowerCAmelCase_      :				Dict      =       """running_var"""
								elif "num_batches_tracked" in name:
										lowerCAmelCase_      :				Union[str, Any]      =       """num_batches_tracked"""
								else:
										lowerCAmelCase_      :				int      =       None
								set_recursively(lowercase__							,   lowercase__							,   lowercase__							,   lowercase__							,   lowercase__ )
						continue
				if not is_used:
						unused_weights.append(lowercase__ )

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







@torch.no_grad()
def     UpperCamelCase_      (   A__ :       Any							,   A__ :       Tuple							,   A__ :       Any							,   A__ :       Optional[Any]=None							,   A__ :       Optional[int]=None							,   ):
		'''simple docstring'''





		if config_path is not None:
				lowerCAmelCase_      :				int      =       EncodecConfig.from_pretrained(lowercase__ )
		else:
				lowerCAmelCase_      :				Optional[int]      =       EncodecConfig()

		if model_name == "encodec_24khz":
				pass  # config is already correct
		elif model_name == "encodec_32khz":
				lowerCAmelCase_      :				List[Any]      =       [8, 5, 4, 4]
				lowerCAmelCase_      :				Tuple      =       [2.2]
				lowerCAmelCase_      :				Optional[Any]      =       64
				lowerCAmelCase_      :				Optional[Any]      =       3_20_00
				lowerCAmelCase_      :				Tuple      =       20_48
				lowerCAmelCase_      :				List[Any]      =       False
				lowerCAmelCase_      :				List[str]      =       False
				lowerCAmelCase_      :				Any      =       False
		elif model_name == "encodec_48khz":
				lowerCAmelCase_      :				Any      =       [8, 5, 4, 2]
				lowerCAmelCase_      :				Any      =       [3.0, 6.0, 12.0, 24.0]
				lowerCAmelCase_      :				List[str]      =       4_80_00
				lowerCAmelCase_      :				List[str]      =       2
				lowerCAmelCase_      :				List[str]      =       False
				lowerCAmelCase_      :				Dict      =       """time_group_norm"""
				lowerCAmelCase_      :				Tuple      =       True
				lowerCAmelCase_      :				Union[str, Any]      =       1.0
				lowerCAmelCase_      :				int      =       0.01
		else:
				raise ValueError(f'Unknown model name: {model_name}' )

		lowerCAmelCase_      :				List[Any]      =       EncodecModel(lowercase__ )

		lowerCAmelCase_      :				Optional[Any]      =       EncodecFeatureExtractor(
		    feature_size=config.audio_channels							,   sampling_rate=config.sampling_rate							,   chunk_length_s=config.chunk_length_s							,   overlap=config.overlap							,   )
		feature_extractor.save_pretrained(lowercase__ )

		lowerCAmelCase_      :				Optional[Any]      =       torch.load(lowercase__ )
		if "best_state" in original_checkpoint:
				# we might have a training state saved, in which case discard the yaml results and just retain the weights
				lowerCAmelCase_      :				Tuple      =       original_checkpoint["""best_state"""]
		recursively_load_weights(lowercase__							,   lowercase__							,   lowercase__ )
		model.save_pretrained(lowercase__ )

		if repo_id:
				print("""Pushing to the hub...""" )
				feature_extractor.push_to_hub(lowercase__ )
				model.push_to_hub(lowercase__ )


if __name__ == "__main__":
		__A :  Dict    							=				argparse.ArgumentParser()
		parser.add_argument(
		    "--model",
		    default="encodec_24khz",
		    type=str,
		    help="The model to convert. Should be one of \'encodec_24khz\', \'encodec_32khz\', \'encodec_48khz\'.",
		)
		parser.add_argument("--checkpoint_path", required=True, default=None, type=str, help="Path to original checkpoint")
		parser.add_argument("--config_path", default=None, type=str, help="Path to hf config.json of model to convert")
		parser.add_argument(
		    "--pytorch_dump_folder_path", required=True, default=None, type=str, help="Path to the output PyTorch model."
		)
		parser.add_argument(
		    "--push_to_hub", default=None, type=str, help="Where to upload the converted model on the 🤗 hub."
		)

		__A :  Tuple    							=				parser.parse_args()
		convert_checkpoint(
		    args.model,
		    args.checkpoint_path,
		    args.pytorch_dump_folder_path,
		    args.config_path,
		    args.push_to_hub,
		)







'''simple docstring'''






import inspect
import unittest

import numpy as np

from tests.test_modeling_common import floats_tensor
from transformers import DetrConfig, MaskFormerConfig, SwinConfig, is_torch_available, is_vision_available
from transformers.testing_utils import require_torch, require_torch_multi_gpu, require_vision, slow, torch_device
from transformers.utils import cached_property

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
  import torch

  from transformers import MaskFormerForInstanceSegmentation, MaskFormerModel

  if is_vision_available():
    from transformers import MaskFormerImageProcessor

if is_vision_available():
  from PIL import Image






class        __snake_case	:







    """simple docstring"""



    def __init__(       self   :  Union[str, Any]   ,						lowerCamelCase   :  List[Any]   ,						lowerCamelCase   :  List[Any]=2   ,						lowerCamelCase   :  int=True   ,						lowerCamelCase   :  str=False   ,						lowerCamelCase   :  List[str]=10   ,						lowerCamelCase   :  Dict=3   ,						lowerCamelCase   :  str=32 * 4   ,						lowerCamelCase   :  Tuple=32 * 6   ,						lowerCamelCase   :  int=4   ,						lowerCamelCase   :  Optional[int]=32   ,						)       ->  List[Any]:
      lowerCAmelCase_      :				Tuple      =       parent
      lowerCAmelCase_      :				int      =       batch_size
      lowerCAmelCase_      :				Tuple      =       is_training
      lowerCAmelCase_      :				str      =       use_auxiliary_loss
      lowerCAmelCase_      :				Optional[Any]      =       num_queries
      lowerCAmelCase_      :				List[str]      =       num_channels
      lowerCAmelCase_      :				Optional[Any]      =       min_size
      lowerCAmelCase_      :				Dict      =       max_size
      lowerCAmelCase_      :				List[Any]      =       num_labels
      lowerCAmelCase_      :				List[Any]      =       mask_feature_size

    def __lowercase       (       self   :  str		)       ->  List[Any]:
      lowerCAmelCase_      :				Union[str, Any]      =       floats_tensor([self.batch_size, self.num_channels, self.min_size, self.max_size]		).to(
          lowerCamelCase		)

      lowerCAmelCase_      :				int      =       torch.ones([self.batch_size, self.min_size, self.max_size]   ,						device=lowerCamelCase		)

      lowerCAmelCase_      :				List[Any]      =       (
          torch.rand([self.batch_size, self.num_labels, self.min_size, self.max_size]   ,						device=lowerCamelCase		) > 0.5
      ).float()
      lowerCAmelCase_      :				Union[str, Any]      =       (torch.rand((self.batch_size, self.num_labels)   ,						device=lowerCamelCase		) > 0.5).long()

      lowerCAmelCase_      :				List[Any]      =       self.get_config()
      return config, pixel_values, pixel_mask, mask_labels, class_labels

    def __lowercase       (       self   :  Optional[int]		)       ->  Optional[int]:
      return MaskFormerConfig.from_backbone_and_decoder_configs(
          backbone_config=SwinConfig(
              depths=[1, 1, 1, 1]   ,						)   ,						decoder_config=DetrConfig(
              decoder_ffn_dim=1_28   ,						num_queries=self.num_queries   ,						decoder_attention_heads=2   ,						d_model=self.mask_feature_size   ,						)   ,						mask_feature_size=self.mask_feature_size   ,						fpn_feature_size=self.mask_feature_size   ,						num_channels=self.num_channels   ,						num_labels=self.num_labels   ,						)

    def __lowercase       (       self   :  Tuple		)       ->  Optional[Any]:
      lowerCAmelCase_,							lowerCAmelCase_,							lowerCAmelCase_,							lowerCAmelCase_,							lowerCAmelCase_      :				List[Any]      =       self.prepare_config_and_inputs()
      lowerCAmelCase_      :				Union[str, Any]      =       {"""pixel_values""": pixel_values, """pixel_mask""": pixel_mask}
      return config, inputs_dict

    def __lowercase       (       self   :  Tuple   ,						lowerCamelCase   :  Tuple   ,						lowerCamelCase   :  Optional[Any]		)       ->  Union[str, Any]:
      lowerCAmelCase_      :				Any      =       output.encoder_hidden_states
      lowerCAmelCase_      :				Dict      =       output.pixel_decoder_hidden_states
      lowerCAmelCase_      :				str      =       output.transformer_decoder_hidden_states

      self.parent.assertTrue(len(lowerCamelCase		)   ,						len(config.backbone_config.depths		)		)
      self.parent.assertTrue(len(lowerCamelCase		)   ,						len(config.backbone_config.depths		)		)
      self.parent.assertTrue(len(lowerCamelCase		)   ,						config.decoder_config.decoder_layers		)

    def __lowercase       (       self   :  List[Any]   ,						lowerCamelCase   :  Any   ,						lowerCamelCase   :  int   ,						lowerCamelCase   :  Optional[int]   ,						lowerCamelCase   :  Tuple=False		)       ->  List[Any]:
      with torch.no_grad():
        lowerCAmelCase_      :				int      =       MaskFormerModel(config=lowerCamelCase		)
        model.to(lowerCamelCase		)
        model.eval()

        lowerCAmelCase_      :				List[Any]      =       model(pixel_values=lowerCamelCase   ,						pixel_mask=lowerCamelCase		)
        lowerCAmelCase_      :				Union[str, Any]      =       model(lowerCamelCase   ,						output_hidden_states=lowerCamelCase		)
      # the correct shape of output.transformer_decoder_hidden_states ensure the correcteness of the
      # encoder and pixel decoder
      self.parent.assertEqual(
          output.transformer_decoder_last_hidden_state.shape   ,						(self.batch_size, self.num_queries, self.mask_feature_size)   ,						)
      # let's ensure the other two hidden state exists
      self.parent.assertTrue(output.pixel_decoder_last_hidden_state is not None		)
      self.parent.assertTrue(output.encoder_last_hidden_state is not None		)

      if output_hidden_states:
        self.check_output_hidden_state(lowerCamelCase   ,						lowerCamelCase		)





    def __lowercase       (       self   :  Any   ,						lowerCamelCase   :  Dict   ,						lowerCamelCase   :  Any   ,						lowerCamelCase   :  Optional[int]   ,						lowerCamelCase   :  Tuple   ,						lowerCamelCase   :  Optional[Any]		)       ->  Dict:
      lowerCAmelCase_      :				Tuple      =       MaskFormerForInstanceSegmentation(config=lowerCamelCase		)
      model.to(lowerCamelCase		)
      model.eval()

      def comm_check_on_output(lowerCamelCase   :  Optional[Any]		):
        # let's still check that all the required stuff is there
        self.parent.assertTrue(result.transformer_decoder_last_hidden_state is not None		)
        self.parent.assertTrue(result.pixel_decoder_last_hidden_state is not None		)
        self.parent.assertTrue(result.encoder_last_hidden_state is not None		)
        # okay, now we need to check the logits shape
        # due to the encoder compression, masks have a //4 spatial size
        self.parent.assertEqual(
            result.masks_queries_logits.shape   ,						(self.batch_size, self.num_queries, self.min_size // 4, self.max_size // 4)   ,						)
        # + 1 for null class
        self.parent.assertEqual(
            result.class_queries_logits.shape   ,						(self.batch_size, self.num_queries, self.num_labels + 1)		)

      with torch.no_grad():
        lowerCAmelCase_      :				Dict      =       model(pixel_values=lowerCamelCase   ,						pixel_mask=lowerCamelCase		)
        lowerCAmelCase_      :				List[str]      =       model(lowerCamelCase		)

        comm_check_on_output(lowerCamelCase		)

        lowerCAmelCase_      :				Union[str, Any]      =       model(
            pixel_values=lowerCamelCase   ,						pixel_mask=lowerCamelCase   ,						mask_labels=lowerCamelCase   ,						class_labels=lowerCamelCase		)

      comm_check_on_output(lowerCamelCase		)

      self.parent.assertTrue(result.loss is not None		)
      self.parent.assertEqual(result.loss.shape   ,						torch.Size([1]		)		)







@require_torch
class        __snake_case	(       _SCREAMING_SNAKE_CASE     ,_SCREAMING_SNAKE_CASE     ,unittest.TestCase):







    """simple docstring"""



    lowercase       			= (MaskFormerModel, MaskFormerForInstanceSegmentation) if is_torch_available() else ()
    lowercase       			= (
        {'feature-extraction': MaskFormerModel, 'image-segmentation': MaskFormerForInstanceSegmentation}
        if is_torch_available()
        else {}
    )

    lowercase       			= False
    lowercase       			= False
    lowercase       			= False
    lowercase       			= False

    def __lowercase       (       self   :  Dict		)       ->  Optional[Any]:
      lowerCAmelCase_      :				Tuple      =       MaskFormerModelTester(self		)
      lowerCAmelCase_      :				str      =       ConfigTester(self   ,						config_class=lowerCamelCase   ,						has_text_modality=lowerCamelCase		)

    def __lowercase       (       self   :  List[Any]		)       ->  Optional[Any]:
      self.config_tester.run_common_tests()

    def __lowercase       (       self   :  Tuple		)       ->  List[Any]:
      lowerCAmelCase_,							lowerCAmelCase_      :				List[str]      =       self.model_tester.prepare_config_and_inputs_for_common()
      self.model_tester.create_and_check_maskformer_model(lowerCamelCase   ,						**lowerCamelCase   ,						output_hidden_states=lowerCamelCase		)

    def __lowercase       (       self   :  Union[str, Any]		)       ->  str:
      lowerCAmelCase_      :				Dict      =       self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_maskformer_instance_segmentation_head_model(*lowerCamelCase		)

    @unittest.skip(reason="""MaskFormer does not use inputs_embeds"""		)
    def __lowercase       (       self   :  Any		)       ->  Optional[Any]:
      pass

    @unittest.skip(reason="""MaskFormer does not have a get_input_embeddings method"""		)
    def __lowercase       (       self   :  Union[str, Any]		)       ->  Optional[Any]:
      pass

    @unittest.skip(reason="""MaskFormer is not a generative model"""		)
    def __lowercase       (       self   :  Any		)       ->  str:
      pass

    @unittest.skip(reason="""MaskFormer does not use token embeddings"""		)
    def __lowercase       (       self   :  List[Any]		)       ->  Any:
      pass

    @require_torch_multi_gpu
    @unittest.skip(
        reason="""MaskFormer has some layers using `add_module` which doesn't work well with `nn.DataParallel`"""		)
    def __lowercase       (       self   :  List[str]		)       ->  Optional[int]:
      pass

    @unittest.skip("""Will be fixed soon by reducing the size of the model used for common tests."""		)
    def __lowercase       (       self   :  str		)       ->  Optional[int]:
      pass

    def __lowercase       (       self   :  Tuple		)       ->  List[str]:
      lowerCAmelCase_,							lowerCAmelCase_      :				Any      =       self.model_tester.prepare_config_and_inputs_for_common()

      for model_class in self.all_model_classes:
        lowerCAmelCase_      :				Optional[Any]      =       model_class(lowerCamelCase		)
        lowerCAmelCase_      :				int      =       inspect.signature(model.forward		)
        # signature.parameters is an OrderedDict => so arg_names order is deterministic
        lowerCAmelCase_      :				List[str]      =       [*signature.parameters.keys()]

        lowerCAmelCase_      :				int      =       ["""pixel_values"""]
        self.assertListEqual(arg_names[:1]   ,						lowerCamelCase		)

    @slow
    def __lowercase       (       self   :  Optional[int]		)       ->  List[str]:
      for model_name in ["facebook/maskformer-swin-small-coco"]:
        lowerCAmelCase_      :				List[str]      =       MaskFormerModel.from_pretrained(lowerCamelCase		)
        self.assertIsNotNone(lowerCamelCase		)

    def __lowercase       (       self   :  Dict		)       ->  int:
      lowerCAmelCase_      :				Any      =       (self.model_tester.min_size,) * 2
      lowerCAmelCase_      :				Dict      =       {
          """pixel_values""": torch.randn((2, 3, *size)   ,						device=lowerCamelCase		),
          """mask_labels""": torch.randn((2, 10, *size)   ,						device=lowerCamelCase		),
          """class_labels""": torch.zeros(2   ,						10   ,						device=lowerCamelCase		).long(),
      }

      lowerCAmelCase_      :				Dict      =       MaskFormerForInstanceSegmentation(MaskFormerConfig()		).to(lowerCamelCase		)
      lowerCAmelCase_      :				str      =       model(**lowerCamelCase		)
      self.assertTrue(outputs.loss is not None		)

    def __lowercase       (       self   :  int		)       ->  int:
      lowerCAmelCase_,							lowerCAmelCase_      :				Union[str, Any]      =       self.model_tester.prepare_config_and_inputs_for_common()
      self.model_tester.create_and_check_maskformer_model(lowerCamelCase   ,						**lowerCamelCase   ,						output_hidden_states=lowerCamelCase		)

    def __lowercase       (       self   :  Dict		)       ->  List[str]:
      lowerCAmelCase_,							lowerCAmelCase_      :				Union[str, Any]      =       self.model_tester.prepare_config_and_inputs_for_common()

      for model_class in self.all_model_classes:
        lowerCAmelCase_      :				List[str]      =       model_class(lowerCamelCase		).to(lowerCamelCase		)
        lowerCAmelCase_      :				Tuple      =       model(**lowerCamelCase   ,						output_attentions=lowerCamelCase		)
        self.assertTrue(outputs.attentions is not None		)

    def __lowercase       (       self   :  int		)       ->  Union[str, Any]:
      if not self.model_tester.is_training:
        return
      # only MaskFormerForInstanceSegmentation has the loss
      lowerCAmelCase_      :				Optional[int]      =       self.all_model_classes[1]
      lowerCAmelCase_,							lowerCAmelCase_,							lowerCAmelCase_,							lowerCAmelCase_,							lowerCAmelCase_      :				str      =       self.model_tester.prepare_config_and_inputs()

      lowerCAmelCase_      :				Union[str, Any]      =       model_class(lowerCamelCase		)
      model.to(lowerCamelCase		)
      model.train()

      lowerCAmelCase_      :				str      =       model(lowerCamelCase   ,						mask_labels=lowerCamelCase   ,						class_labels=lowerCamelCase		).loss
      loss.backward()





    def __lowercase       (       self   :  Any		)       ->  Union[str, Any]:
      # only MaskFormerForInstanceSegmentation has the loss
      lowerCAmelCase_      :				List[str]      =       self.all_model_classes[1]
      lowerCAmelCase_,							lowerCAmelCase_,							lowerCAmelCase_,							lowerCAmelCase_,							lowerCAmelCase_      :				Any      =       self.model_tester.prepare_config_and_inputs()
      lowerCAmelCase_      :				List[str]      =       True
      lowerCAmelCase_      :				List[str]      =       True

      lowerCAmelCase_      :				Dict      =       model_class(lowerCamelCase		)
      model.to(lowerCamelCase		)
      model.train()

      lowerCAmelCase_      :				Union[str, Any]      =       model(lowerCamelCase   ,						mask_labels=lowerCamelCase   ,						class_labels=lowerCamelCase		)

      lowerCAmelCase_      :				Optional[Any]      =       outputs.encoder_hidden_states[0]
      encoder_hidden_states.retain_grad()

      lowerCAmelCase_      :				List[str]      =       outputs.pixel_decoder_hidden_states[0]
      pixel_decoder_hidden_states.retain_grad()
      # we requires_grad=True in inputs_embeds (line 2152), the original implementation don't
      lowerCAmelCase_      :				Optional[int]      =       outputs.transformer_decoder_hidden_states[0]
      transformer_decoder_hidden_states.retain_grad()

      lowerCAmelCase_      :				Tuple      =       outputs.attentions[0]
      attentions.retain_grad()

      outputs.loss.backward(retain_graph=lowerCamelCase		)

      self.assertIsNotNone(encoder_hidden_states.grad		)
      self.assertIsNotNone(pixel_decoder_hidden_states.grad		)
      self.assertIsNotNone(transformer_decoder_hidden_states.grad		)
      self.assertIsNotNone(attentions.grad		)







__A :  List[Any]    							=				1E-4







def     UpperCamelCase_      (   ):
  '''simple docstring'''





  lowerCAmelCase_      :				int      =       Image.open("""./tests/fixtures/tests_samples/COCO/000000039769.png""" )
  return image






@require_vision
@slow
class        __snake_case	(       unittest.TestCase):







    """simple docstring"""



    @cached_property
    def __lowercase       (       self   :  Union[str, Any]		)       ->  List[Any]:
      return (
          MaskFormerImageProcessor.from_pretrained("""facebook/maskformer-swin-small-coco"""		)
          if is_vision_available()
          else None
      )

    def __lowercase       (       self   :  List[Any]		)       ->  Union[str, Any]:
      lowerCAmelCase_      :				Optional[int]      =       MaskFormerModel.from_pretrained("""facebook/maskformer-swin-small-coco"""		).to(lowerCamelCase		)
      lowerCAmelCase_      :				Dict      =       self.default_image_processor
      lowerCAmelCase_      :				Optional[int]      =       prepare_img()
      lowerCAmelCase_      :				Dict      =       image_processor(lowerCamelCase   ,						return_tensors="""pt"""		).to(lowerCamelCase		)
      lowerCAmelCase_      :				Union[str, Any]      =       inputs["""pixel_values"""].shape
      # check size is divisible by 32
      self.assertTrue((inputs_shape[-1] % 32) == 0 and (inputs_shape[-2] % 32) == 0		)
      # check size
      self.assertEqual(lowerCamelCase   ,						(1, 3, 8_00, 10_88)		)

      with torch.no_grad():
        lowerCAmelCase_      :				List[str]      =       model(**lowerCamelCase		)

      lowerCAmelCase_      :				Any      =       torch.tensor(
          [[-0.0_482, 0.9_228, 0.4_951], [-0.2_547, 0.8_017, 0.8_527], [-0.0_069, 0.3_385, -0.0_089]]		).to(lowerCamelCase		)
      self.assertTrue(
          torch.allclose(
              outputs.encoder_last_hidden_state[0, 0, :3, :3]   ,						lowerCamelCase   ,						atol=lowerCamelCase		)		)

      lowerCAmelCase_      :				Union[str, Any]      =       torch.tensor(
          [[-0.8_422, -0.8_434, -0.9_718], [-1.0_144, -0.5_565, -0.4_195], [-1.0_038, -0.4_484, -0.1_961]]		).to(lowerCamelCase		)
      self.assertTrue(
          torch.allclose(
              outputs.pixel_decoder_last_hidden_state[0, 0, :3, :3]   ,						lowerCamelCase   ,						atol=lowerCamelCase		)		)

      lowerCAmelCase_      :				Optional[Any]      =       torch.tensor(
          [[0.2_852, -0.0_159, 0.9_735], [0.6_254, 0.1_858, 0.8_529], [-0.0_680, -0.4_116, 1.8_413]]		).to(lowerCamelCase		)
      self.assertTrue(
          torch.allclose(
              outputs.transformer_decoder_last_hidden_state[0, :3, :3]   ,						lowerCamelCase   ,						atol=lowerCamelCase		)		)

    def __lowercase       (       self   :  Optional[int]		)       ->  Optional[int]:
      lowerCAmelCase_      :				Union[str, Any]      =       (
          MaskFormerForInstanceSegmentation.from_pretrained("""facebook/maskformer-swin-small-coco"""		)
          .to(lowerCamelCase		)
          .eval()
      )
      lowerCAmelCase_      :				Union[str, Any]      =       self.default_image_processor
      lowerCAmelCase_      :				Optional[Any]      =       prepare_img()
      lowerCAmelCase_      :				Optional[int]      =       image_processor(lowerCamelCase   ,						return_tensors="""pt"""		).to(lowerCamelCase		)
      lowerCAmelCase_      :				Optional[int]      =       inputs["""pixel_values"""].shape
      # check size is divisible by 32
      self.assertTrue((inputs_shape[-1] % 32) == 0 and (inputs_shape[-2] % 32) == 0		)
      # check size
      self.assertEqual(lowerCamelCase   ,						(1, 3, 8_00, 10_88)		)

      with torch.no_grad():
        lowerCAmelCase_      :				List[str]      =       model(**lowerCamelCase		)
      # masks_queries_logits
      lowerCAmelCase_      :				Dict      =       outputs.masks_queries_logits
      self.assertEqual(
          masks_queries_logits.shape   ,						(1, model.config.decoder_config.num_queries, inputs_shape[-2] // 4, inputs_shape[-1] // 4)   ,						)
      lowerCAmelCase_      :				Any      =       [
          [-1.3_737_124, -1.7_724_937, -1.9_364_233],
          [-1.5_977_281, -1.9_867_939, -2.1_523_695],
          [-1.5_795_398, -1.9_269_832, -2.093_942],
      ]
      lowerCAmelCase_      :				Optional[int]      =       torch.tensor(lowerCamelCase		).to(lowerCamelCase		)
      self.assertTrue(torch.allclose(masks_queries_logits[0, 0, :3, :3]   ,						lowerCamelCase   ,						atol=lowerCamelCase		)		)
      # class_queries_logits
      lowerCAmelCase_      :				Optional[Any]      =       outputs.class_queries_logits
      self.assertEqual(
          class_queries_logits.shape   ,						(1, model.config.decoder_config.num_queries, model.config.num_labels + 1)		)
      lowerCAmelCase_      :				Tuple      =       torch.tensor(
          [
              [1.6_512E00, -5.2_572E00, -3.3_519E00],
              [3.6_169E-02, -5.9_025E00, -2.9_313E00],
              [1.0_766E-04, -7.7_630E00, -5.1_263E00],
          ]		).to(lowerCamelCase		)
      self.assertTrue(torch.allclose(outputs.class_queries_logits[0, :3, :3]   ,						lowerCamelCase   ,						atol=lowerCamelCase		)		)

    def __lowercase       (       self   :  Optional[int]		)       ->  List[str]:
      lowerCAmelCase_      :				List[str]      =       (
          MaskFormerForInstanceSegmentation.from_pretrained("""facebook/maskformer-resnet101-coco-stuff"""		)
          .to(lowerCamelCase		)
          .eval()
      )
      lowerCAmelCase_      :				Tuple      =       self.default_image_processor
      lowerCAmelCase_      :				int      =       prepare_img()
      lowerCAmelCase_      :				Any      =       image_processor(lowerCamelCase   ,						return_tensors="""pt"""		).to(lowerCamelCase		)
      lowerCAmelCase_      :				List[str]      =       inputs["""pixel_values"""].shape
      # check size is divisible by 32
      self.assertTrue((inputs_shape[-1] % 32) == 0 and (inputs_shape[-2] % 32) == 0		)
      # check size
      self.assertEqual(lowerCamelCase   ,						(1, 3, 8_00, 10_88)		)

      with torch.no_grad():
        lowerCAmelCase_      :				Optional[Any]      =       model(**lowerCamelCase		)
      # masks_queries_logits
      lowerCAmelCase_      :				List[Any]      =       outputs.masks_queries_logits
      self.assertEqual(
          masks_queries_logits.shape   ,						(1, model.config.decoder_config.num_queries, inputs_shape[-2] // 4, inputs_shape[-1] // 4)   ,						)
      lowerCAmelCase_      :				Optional[Any]      =       [[-0.9_046, -2.6_366, -4.6_062], [-3.4_179, -5.7_890, -8.8_057], [-4.9_179, -7.6_560, -10.7_711]]
      lowerCAmelCase_      :				Union[str, Any]      =       torch.tensor(lowerCamelCase		).to(lowerCamelCase		)
      self.assertTrue(torch.allclose(masks_queries_logits[0, 0, :3, :3]   ,						lowerCamelCase   ,						atol=lowerCamelCase		)		)
      # class_queries_logits
      lowerCAmelCase_      :				Dict      =       outputs.class_queries_logits
      self.assertEqual(
          class_queries_logits.shape   ,						(1, model.config.decoder_config.num_queries, model.config.num_labels + 1)		)
      lowerCAmelCase_      :				int      =       torch.tensor(
          [[4.7_188, -3.2_585, -2.8_857], [6.6_871, -2.9_181, -1.2_487], [7.2_449, -2.2_764, -2.1_874]]		).to(lowerCamelCase		)
      self.assertTrue(torch.allclose(outputs.class_queries_logits[0, :3, :3]   ,						lowerCamelCase   ,						atol=lowerCamelCase		)		)





    def __lowercase       (       self   :  Union[str, Any]		)       ->  List[str]:
      lowerCAmelCase_      :				Dict      =       (
          MaskFormerForInstanceSegmentation.from_pretrained("""facebook/maskformer-swin-small-coco"""		)
          .to(lowerCamelCase		)
          .eval()
      )
      lowerCAmelCase_      :				List[str]      =       self.default_image_processor

      lowerCAmelCase_      :				int      =       image_processor(
          [np.zeros((3, 8_00, 13_33)		), np.zeros((3, 8_00, 13_33)		)]   ,						segmentation_maps=[np.zeros((3_84, 3_84)		).astype(np.floataa		), np.zeros((3_84, 3_84)		).astype(np.floataa		)]   ,						return_tensors="""pt"""   ,						)

      lowerCAmelCase_      :				List[str]      =       inputs["""pixel_values"""].to(lowerCamelCase		)
      lowerCAmelCase_      :				Tuple      =       [el.to(lowerCamelCase		) for el in inputs["""mask_labels"""]]
      lowerCAmelCase_      :				Union[str, Any]      =       [el.to(lowerCamelCase		) for el in inputs["""class_labels"""]]

      with torch.no_grad():
        lowerCAmelCase_      :				Any      =       model(**lowerCamelCase		)

      self.assertTrue(outputs.loss is not None		)



