



import argparse
import os

import torch

from diffusers import (
    CMStochasticIterativeScheduler,
    ConsistencyModelPipeline,
    UNetaDModel,
)


UpperCAmelCase      :							str							       ={
    """sample_size""": 32,
    """in_channels""": 3,
    """out_channels""": 3,
    """layers_per_block""": 2,
    """num_class_embeds""": 1000,
    """block_out_channels""": [32, 64],
    """attention_head_dim""": 8,
    """down_block_types""": [
        """ResnetDownsampleBlock2D""",
        """AttnDownBlock2D""",
    ],
    """up_block_types""": [
        """AttnUpBlock2D""",
        """ResnetUpsampleBlock2D""",
    ],
    """resnet_time_scale_shift""": """scale_shift""",
    """upsample_type""": """resnet""",
    """downsample_type""": """resnet""",
}

UpperCAmelCase      :							Dict							       ={
    """sample_size""": 64,
    """in_channels""": 3,
    """out_channels""": 3,
    """layers_per_block""": 3,
    """num_class_embeds""": 1000,
    """block_out_channels""": [192, 192 * 2, 192 * 3, 192 * 4],
    """attention_head_dim""": 64,
    """down_block_types""": [
        """ResnetDownsampleBlock2D""",
        """AttnDownBlock2D""",
        """AttnDownBlock2D""",
        """AttnDownBlock2D""",
    ],
    """up_block_types""": [
        """AttnUpBlock2D""",
        """AttnUpBlock2D""",
        """AttnUpBlock2D""",
        """ResnetUpsampleBlock2D""",
    ],
    """resnet_time_scale_shift""": """scale_shift""",
    """upsample_type""": """resnet""",
    """downsample_type""": """resnet""",
}

UpperCAmelCase      :							Dict							       ={
    """sample_size""": 256,
    """in_channels""": 3,
    """out_channels""": 3,
    """layers_per_block""": 2,
    """num_class_embeds""": None,
    """block_out_channels""": [256, 256, 256 * 2, 256 * 2, 256 * 4, 256 * 4],
    """attention_head_dim""": 64,
    """down_block_types""": [
        """ResnetDownsampleBlock2D""",
        """ResnetDownsampleBlock2D""",
        """ResnetDownsampleBlock2D""",
        """AttnDownBlock2D""",
        """AttnDownBlock2D""",
        """AttnDownBlock2D""",
    ],
    """up_block_types""": [
        """AttnUpBlock2D""",
        """AttnUpBlock2D""",
        """AttnUpBlock2D""",
        """ResnetUpsampleBlock2D""",
        """ResnetUpsampleBlock2D""",
        """ResnetUpsampleBlock2D""",
    ],
    """resnet_time_scale_shift""": """default""",
    """upsample_type""": """resnet""",
    """downsample_type""": """resnet""",
}

UpperCAmelCase      :							int							       ={
    """num_train_timesteps""": 40,
    """sigma_min""": 0.002,
    """sigma_max""": 80.0,
}

UpperCAmelCase      :							List[str]							       ={
    """num_train_timesteps""": 201,
    """sigma_min""": 0.002,
    """sigma_max""": 80.0,
}

UpperCAmelCase      :							Union[str, Any]							       ={
    """num_train_timesteps""": 151,
    """sigma_min""": 0.002,
    """sigma_max""": 80.0,
}



def       _lowerCAmelCase			(_lowerCAmelCase):
					if isinstance(__UpperCAmelCase      ,		__UpperCAmelCase):
										return v
					if v.lower() in ("yes", "true", "t", "y", "1"):
										return True
					elif v.lower() in ("no", "false", "f", "n", "0"):
										return False
					else:
										raise argparse.ArgumentTypeError("boolean value expected")



def       _lowerCAmelCase			(_lowerCAmelCase      ,		_lowerCAmelCase      ,		_lowerCAmelCase      ,		_lowerCAmelCase      ,		_lowerCAmelCase=False):
					UpperCamelCase_						     =      checkpoint[f"""{old_prefix}.in_layers.0.weight"""]
					UpperCamelCase_						     =      checkpoint[f"""{old_prefix}.in_layers.0.bias"""]
					UpperCamelCase_						     =      checkpoint[f"""{old_prefix}.in_layers.2.weight"""]
					UpperCamelCase_						     =      checkpoint[f"""{old_prefix}.in_layers.2.bias"""]
					UpperCamelCase_						     =      checkpoint[f"""{old_prefix}.emb_layers.1.weight"""]
					UpperCamelCase_						     =      checkpoint[f"""{old_prefix}.emb_layers.1.bias"""]
					UpperCamelCase_						     =      checkpoint[f"""{old_prefix}.out_layers.0.weight"""]
					UpperCamelCase_						     =      checkpoint[f"""{old_prefix}.out_layers.0.bias"""]
					UpperCamelCase_						     =      checkpoint[f"""{old_prefix}.out_layers.3.weight"""]
					UpperCamelCase_						     =      checkpoint[f"""{old_prefix}.out_layers.3.bias"""]

					if has_skip:
										UpperCamelCase_						     =      checkpoint[f"""{old_prefix}.skip_connection.weight"""]
										UpperCamelCase_						     =      checkpoint[f"""{old_prefix}.skip_connection.bias"""]

					return new_checkpoint



def       _lowerCAmelCase			(_lowerCAmelCase      ,		_lowerCAmelCase      ,		_lowerCAmelCase      ,		_lowerCAmelCase      ,		_lowerCAmelCase=None):
					UpperCamelCase_						     =      checkpoint[f"""{old_prefix}.qkv.weight"""].chunk(3      ,		dim=0)
					UpperCamelCase_						     =      checkpoint[f"""{old_prefix}.qkv.bias"""].chunk(3      ,		dim=0)

					UpperCamelCase_						     =      checkpoint[f"""{old_prefix}.norm.weight"""]
					UpperCamelCase_						     =      checkpoint[f"""{old_prefix}.norm.bias"""]

					UpperCamelCase_						     =      weight_q.squeeze(-1).squeeze(-1)
					UpperCamelCase_						     =      bias_q.squeeze(-1).squeeze(-1)
					UpperCamelCase_						     =      weight_k.squeeze(-1).squeeze(-1)
					UpperCamelCase_						     =      bias_k.squeeze(-1).squeeze(-1)
					UpperCamelCase_						     =      weight_v.squeeze(-1).squeeze(-1)
					UpperCamelCase_						     =      bias_v.squeeze(-1).squeeze(-1)

					UpperCamelCase_						     =      (
					    checkpoint[f"""{old_prefix}.proj_out.weight"""].squeeze(-1).squeeze(-1)
					)
					UpperCamelCase_						     =      checkpoint[f"""{old_prefix}.proj_out.bias"""].squeeze(-1).squeeze(-1)

					return new_checkpoint



def       _lowerCAmelCase			(_lowerCAmelCase      ,		_lowerCAmelCase):
					UpperCamelCase_						     =      torch.load(__UpperCAmelCase      ,		map_location="cpu")
					UpperCamelCase_						     =      {}

					UpperCamelCase_						     =      checkpoint['''time_embed.0.weight''']
					UpperCamelCase_						     =      checkpoint['''time_embed.0.bias''']
					UpperCamelCase_						     =      checkpoint['''time_embed.2.weight''']
					UpperCamelCase_						     =      checkpoint['''time_embed.2.bias''']

					if unet_config["num_class_embeds"] is not None:
										UpperCamelCase_						     =      checkpoint['''label_emb.weight''']

					UpperCamelCase_						     =      checkpoint['''input_blocks.0.0.weight''']
					UpperCamelCase_						     =      checkpoint['''input_blocks.0.0.bias''']

					UpperCamelCase_						     =      unet_config['''down_block_types''']
					UpperCamelCase_						     =      unet_config['''layers_per_block''']
					UpperCamelCase_						     =      unet_config['''attention_head_dim''']
					UpperCamelCase_						     =      unet_config['''block_out_channels''']
					UpperCamelCase_						     =      1
					UpperCamelCase_						     =      channels_list[0]

					for i, layer_type in enumerate(__UpperCAmelCase):
										UpperCamelCase_						     =      channels_list[i]
										UpperCamelCase_						     =      current_channels != prev_channels
										if layer_type == "ResnetDownsampleBlock2D":
															for j in range(__UpperCAmelCase):
																				UpperCamelCase_						     =      f"""down_blocks.{i}.resnets.{j}"""
																				UpperCamelCase_						     =      f"""input_blocks.{current_layer}.0"""
																				UpperCamelCase_						     =      True if j == 0 and downsample_block_has_skip else False
																				UpperCamelCase_						     =      convert_resnet(__UpperCAmelCase      ,		__UpperCAmelCase      ,		__UpperCAmelCase      ,		__UpperCAmelCase      ,		has_skip=__UpperCAmelCase)
																				current_layer += 1

										elif layer_type == "AttnDownBlock2D":
															for j in range(__UpperCAmelCase):
																				UpperCamelCase_						     =      f"""down_blocks.{i}.resnets.{j}"""
																				UpperCamelCase_						     =      f"""input_blocks.{current_layer}.0"""
																				UpperCamelCase_						     =      True if j == 0 and downsample_block_has_skip else False
																				UpperCamelCase_						     =      convert_resnet(__UpperCAmelCase      ,		__UpperCAmelCase      ,		__UpperCAmelCase      ,		__UpperCAmelCase      ,		has_skip=__UpperCAmelCase)
																				UpperCamelCase_						     =      f"""down_blocks.{i}.attentions.{j}"""
																				UpperCamelCase_						     =      f"""input_blocks.{current_layer}.1"""
																				UpperCamelCase_						     =      convert_attention(
																				    __UpperCAmelCase      ,		__UpperCAmelCase      ,		__UpperCAmelCase      ,		__UpperCAmelCase      ,		__UpperCAmelCase)
																				current_layer += 1

										if i != len(__UpperCAmelCase) - 1:
															UpperCamelCase_						     =      f"""down_blocks.{i}.downsamplers.0"""
															UpperCamelCase_						     =      f"""input_blocks.{current_layer}.0"""
															UpperCamelCase_						     =      convert_resnet(__UpperCAmelCase      ,		__UpperCAmelCase      ,		__UpperCAmelCase      ,		__UpperCAmelCase)
															current_layer += 1

										UpperCamelCase_						     =      current_channels

					# hardcoded the mid-block for now
					UpperCamelCase_						     =      '''mid_block.resnets.0'''
					UpperCamelCase_						     =      '''middle_block.0'''
					UpperCamelCase_						     =      convert_resnet(__UpperCAmelCase      ,		__UpperCAmelCase      ,		__UpperCAmelCase      ,		__UpperCAmelCase)
					UpperCamelCase_						     =      '''mid_block.attentions.0'''
					UpperCamelCase_						     =      '''middle_block.1'''
					UpperCamelCase_						     =      convert_attention(__UpperCAmelCase      ,		__UpperCAmelCase      ,		__UpperCAmelCase      ,		__UpperCAmelCase      ,		__UpperCAmelCase)
					UpperCamelCase_						     =      '''mid_block.resnets.1'''
					UpperCamelCase_						     =      '''middle_block.2'''
					UpperCamelCase_						     =      convert_resnet(__UpperCAmelCase      ,		__UpperCAmelCase      ,		__UpperCAmelCase      ,		__UpperCAmelCase)

					UpperCamelCase_						     =      0
					UpperCamelCase_						     =      unet_config['''up_block_types''']

					for i, layer_type in enumerate(__UpperCAmelCase):
										if layer_type == "ResnetUpsampleBlock2D":
															for j in range(layers_per_block + 1):
																				UpperCamelCase_						     =      f"""up_blocks.{i}.resnets.{j}"""
																				UpperCamelCase_						     =      f"""output_blocks.{current_layer}.0"""
																				UpperCamelCase_						     =      convert_resnet(__UpperCAmelCase      ,		__UpperCAmelCase      ,		__UpperCAmelCase      ,		__UpperCAmelCase      ,		has_skip=__UpperCAmelCase)
																				current_layer += 1

															if i != len(__UpperCAmelCase) - 1:
																				UpperCamelCase_						     =      f"""up_blocks.{i}.upsamplers.0"""
																				UpperCamelCase_						     =      f"""output_blocks.{current_layer-1}.1"""
																				UpperCamelCase_						     =      convert_resnet(__UpperCAmelCase      ,		__UpperCAmelCase      ,		__UpperCAmelCase      ,		__UpperCAmelCase)
										elif layer_type == "AttnUpBlock2D":
															for j in range(layers_per_block + 1):
																				UpperCamelCase_						     =      f"""up_blocks.{i}.resnets.{j}"""
																				UpperCamelCase_						     =      f"""output_blocks.{current_layer}.0"""
																				UpperCamelCase_						     =      convert_resnet(__UpperCAmelCase      ,		__UpperCAmelCase      ,		__UpperCAmelCase      ,		__UpperCAmelCase      ,		has_skip=__UpperCAmelCase)
																				UpperCamelCase_						     =      f"""up_blocks.{i}.attentions.{j}"""
																				UpperCamelCase_						     =      f"""output_blocks.{current_layer}.1"""
																				UpperCamelCase_						     =      convert_attention(
																				    __UpperCAmelCase      ,		__UpperCAmelCase      ,		__UpperCAmelCase      ,		__UpperCAmelCase      ,		__UpperCAmelCase)
																				current_layer += 1

															if i != len(__UpperCAmelCase) - 1:
																				UpperCamelCase_						     =      f"""up_blocks.{i}.upsamplers.0"""
																				UpperCamelCase_						     =      f"""output_blocks.{current_layer-1}.2"""
																				UpperCamelCase_						     =      convert_resnet(__UpperCAmelCase      ,		__UpperCAmelCase      ,		__UpperCAmelCase      ,		__UpperCAmelCase)

					UpperCamelCase_						     =      checkpoint['''out.0.weight''']
					UpperCamelCase_						     =      checkpoint['''out.0.bias''']
					UpperCamelCase_						     =      checkpoint['''out.2.weight''']
					UpperCamelCase_						     =      checkpoint['''out.2.bias''']

					return new_checkpoint


if __name__ == "__main__":
	UpperCAmelCase      :							List[str]							       =argparse.ArgumentParser()

	parser.add_argument("""--unet_path""", default=None, type=str, required=True, help="""Path to the unet.pt to convert.""")
	parser.add_argument(
	    """--dump_path""", default=None, type=str, required=True, help="""Path to output the converted UNet model."""
	)
	parser.add_argument("""--class_cond""", default=True, type=str, help="""Whether the model is class-conditional.""")

	UpperCAmelCase      :							List[Any]							       =parser.parse_args()
	UpperCAmelCase      :							Any							       =strabool(args.class_cond)

	UpperCAmelCase      :							Optional[Any]							       =os.path.basename(args.unet_path)
	print(F"Checkpoint: {ckpt_name}")

	# Get U-Net config
	if "imagenet64" in ckpt_name:
		UpperCAmelCase      :							List[str]							       =IMAGENET_64_UNET_CONFIG
	elif "256" in ckpt_name and (("bedroom" in ckpt_name) or ("cat" in ckpt_name)):
		UpperCAmelCase      :							Union[str, Any]							       =LSUN_256_UNET_CONFIG
	elif "test" in ckpt_name:
		UpperCAmelCase      :							Dict							       =TEST_UNET_CONFIG
	else:
		raise ValueError(F"Checkpoint type {ckpt_name} is not currently supported.")

	if not args.class_cond:
		UpperCAmelCase      :							Optional[int]							       =None

	UpperCAmelCase      :							Union[str, Any]							       =con_pt_to_diffuser(args.unet_path, unet_config)

	UpperCAmelCase      :							List[str]							       =UNetaDModel(**unet_config)
	image_unet.load_state_dict(converted_unet_ckpt)

	# Get scheduler config
	if "cd" in ckpt_name or "test" in ckpt_name:
		UpperCAmelCase      :							Tuple							       =CD_SCHEDULER_CONFIG
	elif "ct" in ckpt_name and "imagenet64" in ckpt_name:
		UpperCAmelCase      :							Dict							       =CT_IMAGENET_64_SCHEDULER_CONFIG
	elif "ct" in ckpt_name and "256" in ckpt_name and (("bedroom" in ckpt_name) or ("cat" in ckpt_name)):
		UpperCAmelCase      :							Dict							       =CT_LSUN_256_SCHEDULER_CONFIG
	else:
		raise ValueError(F"Checkpoint type {ckpt_name} is not currently supported.")

	UpperCAmelCase      :							int							       =CMStochasticIterativeScheduler(**scheduler_config)

	UpperCAmelCase      :							int							       =ConsistencyModelPipeline(unet=image_unet, scheduler=cm_scheduler)
	consistency_model.save_pretrained(args.dump_path)

"""simple docstring"""






from __future__ import annotations

import inspect
import unittest

import numpy as np

from transformers import ResNetConfig
from transformers.testing_utils import require_tf, require_vision, slow
from transformers.utils import cached_property, is_tf_available, is_vision_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
       import tensorflow as tf

       from transformers import TFResNetForImageClassification, TFResNetModel
       from transformers.models.resnet.modeling_tf_resnet import TF_RESNET_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
       from PIL import Image

       from transformers import AutoImageProcessor


class     UpperCAmelCase						:



    """simple docstring"""

    def __init__(   self    ,			_UpperCAmelCase    ,			_UpperCAmelCase=3    ,			_UpperCAmelCase=32    ,			_UpperCAmelCase=3    ,			_UpperCAmelCase=10    ,			_UpperCAmelCase=[10, 20, 30, 40]    ,			_UpperCAmelCase=[1, 1, 2, 1]    ,			_UpperCAmelCase=True    ,			_UpperCAmelCase=True    ,			_UpperCAmelCase="relu"    ,			_UpperCAmelCase=3    ,			_UpperCAmelCase=None    ,			):
       lowercase__:   Optional[Any]					=      parent
       lowercase__:   Union[str, Any]					=      batch_size
       lowercase__:   int					=      image_size
       lowercase__:   Optional[Any]					=      num_channels
       lowercase__:   Optional[int]					=      embeddings_size
       lowercase__:   Dict					=      hidden_sizes
       lowercase__:   Union[str, Any]					=      depths
       lowercase__:   str					=      is_training
       lowercase__:   Optional[int]					=      use_labels
       lowercase__:   List[str]					=      hidden_act
       lowercase__:   Dict					=      num_labels
       lowercase__:   Any					=      scope
       lowercase__:   Optional[Any]					=      len(_UpperCAmelCase							)






    def _snake_case			(   self							):
       lowercase__:   Union[str, Any]					=      floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]							)

       lowercase__:   List[Any]					=      None
       if self.use_labels:
          lowercase__:   Any					=      ids_tensor([self.batch_size]    ,			self.num_labels							)

       lowercase__:   Optional[int]					=      self.get_config()

       return config, pixel_values, labels






    def _snake_case			(   self							):
       return ResNetConfig(
           num_channels=self.num_channels    ,			embeddings_size=self.embeddings_size    ,			hidden_sizes=self.hidden_sizes    ,			depths=self.depths    ,			hidden_act=self.hidden_act    ,			num_labels=self.num_labels    ,			image_size=self.image_size    ,			)






    def _snake_case			(   self    ,			_UpperCAmelCase    ,			_UpperCAmelCase    ,			_UpperCAmelCase							):
       lowercase__:   Optional[Any]					=      TFResNetModel(config=_UpperCAmelCase							)
       lowercase__:   Dict					=      model(_UpperCAmelCase							)
       # expected last hidden states: B, C, H // 32, W // 32
       self.parent.assertEqual(
           result.last_hidden_state.shape    ,			(self.batch_size, self.hidden_sizes[-1], self.image_size // 32, self.image_size // 32)    ,			)






    def _snake_case			(   self    ,			_UpperCAmelCase    ,			_UpperCAmelCase    ,			_UpperCAmelCase							):
       lowercase__:   str					=      self.num_labels
       lowercase__:   int					=      TFResNetForImageClassification(_UpperCAmelCase							)
       lowercase__:   Optional[Any]					=      model(_UpperCAmelCase    ,			labels=_UpperCAmelCase							)
       self.parent.assertEqual(result.logits.shape    ,			(self.batch_size, self.num_labels)							)






    def _snake_case			(   self							):
       lowercase__:   int					=      self.prepare_config_and_inputs()
       lowercase__, lowercase__, lowercase__:   Optional[Any]					=      config_and_inputs
       lowercase__:   Union[str, Any]					=      {'''pixel_values''': pixel_values}
       return config, inputs_dict



@require_tf
class     UpperCAmelCase						(_UpperCAmelCase    ,_UpperCAmelCase    ,unittest.TestCase  ):



    """simple docstring"""

    _UpperCAmelCase	:Any						    = (TFResNetModel, TFResNetForImageClassification) if is_tf_available() else ()
    _UpperCAmelCase	:List[str]						    = (
        {"feature-extraction": TFResNetModel, "image-classification": TFResNetForImageClassification}
        if is_tf_available()
        else {}
    )

    _UpperCAmelCase	:Any						    = False
    _UpperCAmelCase	:List[str]						    = False
    _UpperCAmelCase	:Optional[Any]						    = False
    _UpperCAmelCase	:Tuple						    = False
    _UpperCAmelCase	:List[Any]						    = False






    def _snake_case			(   self							):
       lowercase__:   Union[str, Any]					=      TFResNetModelTester(self							)
       lowercase__:   Optional[int]					=      ConfigTester(self    ,			config_class=_UpperCAmelCase    ,			has_text_modality=_UpperCAmelCase							)






    def _snake_case			(   self							):
       self.create_and_test_config_common_properties()
       self.config_tester.create_and_test_config_to_json_string()
       self.config_tester.create_and_test_config_to_json_file()
       self.config_tester.create_and_test_config_from_and_save_pretrained()
       self.config_tester.create_and_test_config_with_num_labels()
       self.config_tester.check_config_can_be_init_without_params()
       self.config_tester.check_config_arguments_init()






    def _snake_case			(   self							):
       return






    @unittest.skip(reason='''ResNet does not use inputs_embeds'''							)
    def _snake_case			(   self							):
       pass






    @unittest.skip(reason='''ResNet does not support input and output embeddings'''							)
    def _snake_case			(   self							):
       pass






    def _snake_case			(   self							):
       lowercase__, lowercase__:   int					=      self.model_tester.prepare_config_and_inputs_for_common()

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

          lowercase__:   Optional[Any]					=      ['''pixel_values''']
          self.assertListEqual(arg_names[:1]    ,			_UpperCAmelCase							)






    def _snake_case			(   self							):
       lowercase__:   Any					=      self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_model(*_UpperCAmelCase							)






    def _snake_case			(   self							):
       def check_hidden_states_output(_UpperCAmelCase    ,			_UpperCAmelCase    ,			_UpperCAmelCase							):
          lowercase__:   Union[str, Any]					=      model_class(_UpperCAmelCase							)
          lowercase__:   List[Any]					=      model(**self._prepare_for_class(_UpperCAmelCase    ,			_UpperCAmelCase							)							)

          lowercase__:   Union[str, Any]					=      outputs.encoder_hidden_states if config.is_encoder_decoder else outputs.hidden_states

          lowercase__:   Tuple					=      self.model_tester.num_stages
          self.assertEqual(len(_UpperCAmelCase							)    ,			expected_num_stages + 1							)

          # ResNet's feature maps are of shape (batch_size, num_channels, height, width)
          self.assertListEqual(
              list(hidden_states[0].shape[-2:]							)    ,			[self.model_tester.image_size // 4, self.model_tester.image_size // 4]    ,			)

       lowercase__, lowercase__:   Optional[Any]					=      self.model_tester.prepare_config_and_inputs_for_common()
       lowercase__:   Union[str, Any]					=      ['''basic''', '''bottleneck''']
       for model_class in self.all_model_classes:
          for layer_type in layers_type:
             lowercase__:   Tuple					=      layer_type
             lowercase__:   Any					=      True
             check_hidden_states_output(_UpperCAmelCase    ,			_UpperCAmelCase    ,			_UpperCAmelCase							)

             # check that output_hidden_states also work using config
             del inputs_dict["output_hidden_states"]
             lowercase__:   List[str]					=      True

             check_hidden_states_output(_UpperCAmelCase    ,			_UpperCAmelCase    ,			_UpperCAmelCase							)






    def _snake_case			(   self							):
       lowercase__:   Optional[int]					=      self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_image_classification(*_UpperCAmelCase							)






    @slow
    def _snake_case			(   self							):
       for model_name in TF_RESNET_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
          lowercase__:   Dict					=      TFResNetModel.from_pretrained(_UpperCAmelCase							)
          self.assertIsNotNone(_UpperCAmelCase							)



def SCREAMING_SNAKE_CASE__			(   )  ->  List[Any]:
   lowercase__:   List[Any]					=      Image.open('''./tests/fixtures/tests_samples/COCO/000000039769.png''' )
   return image




@require_tf
@require_vision
class     UpperCAmelCase						(unittest.TestCase  ):



    """simple docstring"""

    @cached_property
    def _snake_case			(   self							):
       return (
           AutoImageProcessor.from_pretrained(TF_RESNET_PRETRAINED_MODEL_ARCHIVE_LIST[0]							)
           if is_vision_available()
           else None
       )






    @slow
    def _snake_case			(   self							):
       lowercase__:   Tuple					=      TFResNetForImageClassification.from_pretrained(TF_RESNET_PRETRAINED_MODEL_ARCHIVE_LIST[0]							)

       lowercase__:   Any					=      self.default_image_processor
       lowercase__:   List[Any]					=      prepare_img()
       lowercase__:   List[Any]					=      image_processor(images=_UpperCAmelCase    ,			return_tensors='''tf'''							)

       # forward pass
       lowercase__:   Dict					=      model(**_UpperCAmelCase							)

       # verify the logits
       lowercase__:   int					=      tf.TensorShape((1, 1000)							)
       self.assertEqual(outputs.logits.shape    ,			_UpperCAmelCase							)

       lowercase__:   Dict					=      tf.constant([-11.1_069, -9.7_877, -8.3_777]							)

       self.assertTrue(np.allclose(outputs.logits[0, :3].numpy()    ,			_UpperCAmelCase    ,			atol=1e-4							)							)
