



import argparse
import os

import torch

from diffusers import (
    CMStochasticIterativeScheduler,
    ConsistencyModelPipeline,
    UNetaDModel,
)


__UpperCAmelCase          =       {
    '''sample_size''': 32,
    '''in_channels''': 3,
    '''out_channels''': 3,
    '''layers_per_block''': 2,
    '''num_class_embeds''': 1_000,
    '''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          =       {
    '''sample_size''': 64,
    '''in_channels''': 3,
    '''out_channels''': 3,
    '''layers_per_block''': 3,
    '''num_class_embeds''': 1_000,
    '''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          =       {
    '''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          =       {
    '''num_train_timesteps''': 40,
    '''sigma_min''': 0.0_02,
    '''sigma_max''': 80.0,
}

__UpperCAmelCase          =       {
    '''num_train_timesteps''': 201,
    '''sigma_min''': 0.0_02,
    '''sigma_max''': 80.0,
}

__UpperCAmelCase          =       {
    '''num_train_timesteps''': 151,
    '''sigma_min''': 0.0_02,
    '''sigma_max''': 80.0,
}
def 		UpperCamelCase	(   snake_case__       :    Optional[int]	)	->     Tuple:
   if isinstance(snake_case__       ,     snake_case__	):
      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 		UpperCamelCase	(   snake_case__       :    List[Any]       ,     snake_case__       :    Optional[int]       ,     snake_case__       :    Optional[int]       ,     snake_case__       :    Union[str, Any]       ,     snake_case__       :    Tuple=False	)	->     List[Any]:
   UpperCamelCase    :  Union[str, Any]							=	checkpoint[F"""{old_prefix}.in_layers.0.weight"""]
   UpperCamelCase    :  Union[str, Any]							=	checkpoint[F"""{old_prefix}.in_layers.0.bias"""]
   UpperCamelCase    :  Optional[int]							=	checkpoint[F"""{old_prefix}.in_layers.2.weight"""]
   UpperCamelCase    :  int							=	checkpoint[F"""{old_prefix}.in_layers.2.bias"""]
   UpperCamelCase    :  str							=	checkpoint[F"""{old_prefix}.emb_layers.1.weight"""]
   UpperCamelCase    :  int							=	checkpoint[F"""{old_prefix}.emb_layers.1.bias"""]
   UpperCamelCase    :  Dict							=	checkpoint[F"""{old_prefix}.out_layers.0.weight"""]
   UpperCamelCase    :  int							=	checkpoint[F"""{old_prefix}.out_layers.0.bias"""]
   UpperCamelCase    :  Tuple							=	checkpoint[F"""{old_prefix}.out_layers.3.weight"""]
   UpperCamelCase    :  Optional[int]							=	checkpoint[F"""{old_prefix}.out_layers.3.bias"""]

   if has_skip:
      UpperCamelCase    :  List[Any]							=	checkpoint[F"""{old_prefix}.skip_connection.weight"""]
      UpperCamelCase    :  Any							=	checkpoint[F"""{old_prefix}.skip_connection.bias"""]

   return new_checkpoint
def 		UpperCamelCase	(   snake_case__       :    List[str]       ,     snake_case__       :    List[Any]       ,     snake_case__       :    Optional[Any]       ,     snake_case__       :    int       ,     snake_case__       :    List[Any]=None	)	->     Tuple:
   UpperCamelCase	,     UpperCamelCase	,     UpperCamelCase    :  int							=	checkpoint[F"""{old_prefix}.qkv.weight"""].chunk(3       ,     dim=0	)
   UpperCamelCase	,     UpperCamelCase	,     UpperCamelCase    :  str							=	checkpoint[F"""{old_prefix}.qkv.bias"""].chunk(3       ,     dim=0	)

   UpperCamelCase    :  int							=	checkpoint[F"""{old_prefix}.norm.weight"""]
   UpperCamelCase    :  Union[str, Any]							=	checkpoint[F"""{old_prefix}.norm.bias"""]

   UpperCamelCase    :  Union[str, Any]							=	weight_q.squeeze(-1	).squeeze(-1	)
   UpperCamelCase    :  int							=	bias_q.squeeze(-1	).squeeze(-1	)
   UpperCamelCase    :  str							=	weight_k.squeeze(-1	).squeeze(-1	)
   UpperCamelCase    :  str							=	bias_k.squeeze(-1	).squeeze(-1	)
   UpperCamelCase    :  List[str]							=	weight_v.squeeze(-1	).squeeze(-1	)
   UpperCamelCase    :  Any							=	bias_v.squeeze(-1	).squeeze(-1	)

   UpperCamelCase    :  Union[str, Any]							=	(
       checkpoint[F"""{old_prefix}.proj_out.weight"""].squeeze(-1	).squeeze(-1	)
   )
   UpperCamelCase    :  Any							=	checkpoint[F"""{old_prefix}.proj_out.bias"""].squeeze(-1	).squeeze(-1	)

   return new_checkpoint
def 		UpperCamelCase	(   snake_case__       :    str       ,     snake_case__       :    List[str]	)	->     List[Any]:
   UpperCamelCase    :  List[Any]							=	torch.load(snake_case__       ,     map_location='cpu'	)
   UpperCamelCase    :  Dict							=	{}

   UpperCamelCase    :  List[str]							=	checkpoint['time_embed.0.weight']
   UpperCamelCase    :  Tuple							=	checkpoint['time_embed.0.bias']
   UpperCamelCase    :  str							=	checkpoint['time_embed.2.weight']
   UpperCamelCase    :  Union[str, Any]							=	checkpoint['time_embed.2.bias']

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

   UpperCamelCase    :  List[str]							=	checkpoint['input_blocks.0.0.weight']
   UpperCamelCase    :  Union[str, Any]							=	checkpoint['input_blocks.0.0.bias']

   UpperCamelCase    :  Dict							=	unet_config['down_block_types']
   UpperCamelCase    :  List[Any]							=	unet_config['layers_per_block']
   UpperCamelCase    :  Optional[Any]							=	unet_config['attention_head_dim']
   UpperCamelCase    :  Optional[Any]							=	unet_config['block_out_channels']
   UpperCamelCase    :  Tuple							=	1
   UpperCamelCase    :  Dict							=	channels_list[0]

   for i, layer_type in enumerate(snake_case__	):
      UpperCamelCase    :  Any							=	channels_list[i]
      UpperCamelCase    :  Dict							=	current_channels != prev_channels
      if layer_type == "ResnetDownsampleBlock2D":
         for j in range(snake_case__	):
            UpperCamelCase    :  Optional[Any]							=	F"""down_blocks.{i}.resnets.{j}"""
            UpperCamelCase    :  Tuple							=	F"""input_blocks.{current_layer}.0"""
            UpperCamelCase    :  List[str]							=	True if j == 0 and downsample_block_has_skip else False
            UpperCamelCase    :  List[str]							=	convert_resnet(snake_case__       ,     snake_case__       ,     snake_case__       ,     snake_case__       ,     has_skip=snake_case__	)
            current_layer += 1

      elif layer_type == "AttnDownBlock2D":
         for j in range(snake_case__	):
            UpperCamelCase    :  List[str]							=	F"""down_blocks.{i}.resnets.{j}"""
            UpperCamelCase    :  Optional[Any]							=	F"""input_blocks.{current_layer}.0"""
            UpperCamelCase    :  Union[str, Any]							=	True if j == 0 and downsample_block_has_skip else False
            UpperCamelCase    :  List[Any]							=	convert_resnet(snake_case__       ,     snake_case__       ,     snake_case__       ,     snake_case__       ,     has_skip=snake_case__	)
            UpperCamelCase    :  Optional[int]							=	F"""down_blocks.{i}.attentions.{j}"""
            UpperCamelCase    :  int							=	F"""input_blocks.{current_layer}.1"""
            UpperCamelCase    :  List[str]							=	convert_attention(
                snake_case__       ,     snake_case__       ,     snake_case__       ,     snake_case__       ,     snake_case__	)
            current_layer += 1

      if i != len(snake_case__	) - 1:
         UpperCamelCase    :  Tuple							=	F"""down_blocks.{i}.downsamplers.0"""
         UpperCamelCase    :  int							=	F"""input_blocks.{current_layer}.0"""
         UpperCamelCase    :  str							=	convert_resnet(snake_case__       ,     snake_case__       ,     snake_case__       ,     snake_case__	)
         current_layer += 1

      UpperCamelCase    :  Any							=	current_channels

   # hardcoded the mid-block for now
   UpperCamelCase    :  List[Any]							=	'mid_block.resnets.0'
   UpperCamelCase    :  Dict							=	'middle_block.0'
   UpperCamelCase    :  Optional[int]							=	convert_resnet(snake_case__       ,     snake_case__       ,     snake_case__       ,     snake_case__	)
   UpperCamelCase    :  int							=	'mid_block.attentions.0'
   UpperCamelCase    :  Optional[int]							=	'middle_block.1'
   UpperCamelCase    :  List[Any]							=	convert_attention(snake_case__       ,     snake_case__       ,     snake_case__       ,     snake_case__       ,     snake_case__	)
   UpperCamelCase    :  Optional[Any]							=	'mid_block.resnets.1'
   UpperCamelCase    :  Union[str, Any]							=	'middle_block.2'
   UpperCamelCase    :  str							=	convert_resnet(snake_case__       ,     snake_case__       ,     snake_case__       ,     snake_case__	)

   UpperCamelCase    :  Optional[int]							=	0
   UpperCamelCase    :  str							=	unet_config['up_block_types']

   for i, layer_type in enumerate(snake_case__	):
      if layer_type == "ResnetUpsampleBlock2D":
         for j in range(layers_per_block + 1	):
            UpperCamelCase    :  Tuple							=	F"""up_blocks.{i}.resnets.{j}"""
            UpperCamelCase    :  Tuple							=	F"""output_blocks.{current_layer}.0"""
            UpperCamelCase    :  Any							=	convert_resnet(snake_case__       ,     snake_case__       ,     snake_case__       ,     snake_case__       ,     has_skip=snake_case__	)
            current_layer += 1

         if i != len(snake_case__	) - 1:
            UpperCamelCase    :  Union[str, Any]							=	F"""up_blocks.{i}.upsamplers.0"""
            UpperCamelCase    :  Optional[int]							=	F"""output_blocks.{current_layer-1}.1"""
            UpperCamelCase    :  int							=	convert_resnet(snake_case__       ,     snake_case__       ,     snake_case__       ,     snake_case__	)
      elif layer_type == "AttnUpBlock2D":
         for j in range(layers_per_block + 1	):
            UpperCamelCase    :  Tuple							=	F"""up_blocks.{i}.resnets.{j}"""
            UpperCamelCase    :  Tuple							=	F"""output_blocks.{current_layer}.0"""
            UpperCamelCase    :  int							=	convert_resnet(snake_case__       ,     snake_case__       ,     snake_case__       ,     snake_case__       ,     has_skip=snake_case__	)
            UpperCamelCase    :  int							=	F"""up_blocks.{i}.attentions.{j}"""
            UpperCamelCase    :  List[Any]							=	F"""output_blocks.{current_layer}.1"""
            UpperCamelCase    :  Tuple							=	convert_attention(
                snake_case__       ,     snake_case__       ,     snake_case__       ,     snake_case__       ,     snake_case__	)
            current_layer += 1

         if i != len(snake_case__	) - 1:
            UpperCamelCase    :  List[str]							=	F"""up_blocks.{i}.upsamplers.0"""
            UpperCamelCase    :  Tuple							=	F"""output_blocks.{current_layer-1}.2"""
            UpperCamelCase    :  List[Any]							=	convert_resnet(snake_case__       ,     snake_case__       ,     snake_case__       ,     snake_case__	)

   UpperCamelCase    :  Union[str, Any]							=	checkpoint['out.0.weight']
   UpperCamelCase    :  List[str]							=	checkpoint['out.0.bias']
   UpperCamelCase    :  Optional[Any]							=	checkpoint['out.2.weight']
   UpperCamelCase    :  List[str]							=	checkpoint['out.2.bias']

   return new_checkpoint


if __name__ == "__main__":
    __UpperCAmelCase          =       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          =       parser.parse_args()
    __UpperCAmelCase          =       strabool(args.class_cond)

    __UpperCAmelCase          =       os.path.basename(args.unet_path)
    print(F"""Checkpoint: {ckpt_name}""")

    # Get U-Net config
    if "imagenet64" in ckpt_name:
        __UpperCAmelCase          =       IMAGENET_64_UNET_CONFIG
    elif "256" in ckpt_name and (("bedroom" in ckpt_name) or ("cat" in ckpt_name)):
        __UpperCAmelCase          =       LSUN_256_UNET_CONFIG
    elif "test" in ckpt_name:
        __UpperCAmelCase          =       TEST_UNET_CONFIG
    else:
        raise ValueError(F"""Checkpoint type {ckpt_name} is not currently supported.""")

    if not args.class_cond:
        __UpperCAmelCase          =       None

    __UpperCAmelCase          =       con_pt_to_diffuser(args.unet_path, unet_config)

    __UpperCAmelCase          =       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          =       CD_SCHEDULER_CONFIG
    elif "ct" in ckpt_name and "imagenet64" in ckpt_name:
        __UpperCAmelCase          =       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          =       CT_LSUN_256_SCHEDULER_CONFIG
    else:
        raise ValueError(F"""Checkpoint type {ckpt_name} is not currently supported.""")

    __UpperCAmelCase          =       CMStochasticIterativeScheduler(**scheduler_config)

    __UpperCAmelCase          =       ConsistencyModelPipeline(unet=image_unet, scheduler=cm_scheduler)
    consistency_model.save_pretrained(args.dump_path)








import inspect
import unittest

from transformers import ConvNextConfig
from transformers.testing_utils import require_torch, require_vision, slow, torch_device
from transformers.utils import cached_property, is_torch_available, is_vision_available

from ...test_backbone_common import BackboneTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
    import torch

    from transformers import ConvNextBackbone, ConvNextForImageClassification, ConvNextModel
    from transformers.models.convnext.modeling_convnext import CONVNEXT_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
    from PIL import Image

    from transformers import AutoImageProcessor



class 				lowerCAmelCase_		:
       def __init__(			self,			SCREAMING_SNAKE_CASE_,			SCREAMING_SNAKE_CASE_=13,			SCREAMING_SNAKE_CASE_=32,			SCREAMING_SNAKE_CASE_=3,			SCREAMING_SNAKE_CASE_=4,			SCREAMING_SNAKE_CASE_=[10, 20, 30, 40],			SCREAMING_SNAKE_CASE_=[2, 2, 3, 2],			SCREAMING_SNAKE_CASE_=True,			SCREAMING_SNAKE_CASE_=True,			SCREAMING_SNAKE_CASE_=37,			SCREAMING_SNAKE_CASE_="gelu",			SCREAMING_SNAKE_CASE_=10,			SCREAMING_SNAKE_CASE_=0.02,			SCREAMING_SNAKE_CASE_=["stage2", "stage3", "stage4"],			SCREAMING_SNAKE_CASE_=[2, 3, 4],			SCREAMING_SNAKE_CASE_=None,			)     ->		Optional[int]:
          UpperCamelCase    :  Dict							=	parent
          UpperCamelCase    :  Optional[Any]							=	batch_size
          UpperCamelCase    :  Union[str, Any]							=	image_size
          UpperCamelCase    :  Union[str, Any]							=	num_channels
          UpperCamelCase    :  List[Any]							=	num_stages
          UpperCamelCase    :  Any							=	hidden_sizes
          UpperCamelCase    :  Optional[int]							=	depths
          UpperCamelCase    :  Optional[int]							=	is_training
          UpperCamelCase    :  List[str]							=	use_labels
          UpperCamelCase    :  Dict							=	intermediate_size
          UpperCamelCase    :  List[Any]							=	hidden_act
          UpperCamelCase    :  Union[str, Any]							=	num_labels
          UpperCamelCase    :  str							=	initializer_range
          UpperCamelCase    :  List[str]							=	out_features
          UpperCamelCase    :  List[str]							=	out_indices
          UpperCamelCase    :  str							=	scope
       def 					snake_case_       (			self       )     ->		Optional[Any]:
          UpperCamelCase    :  int							=	floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]       )

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

          UpperCamelCase    :  List[Any]							=	self.get_config()
          return config, pixel_values, labels
       def 					snake_case_       (			self       )     ->		Any:
          return ConvNextConfig(
              num_channels=self.num_channels,			hidden_sizes=self.hidden_sizes,			depths=self.depths,			num_stages=self.num_stages,			hidden_act=self.hidden_act,			is_decoder=SCREAMING_SNAKE_CASE_,			initializer_range=self.initializer_range,			out_features=self.out_features,			out_indices=self.out_indices,			num_labels=self.num_labels,			)
       def 					snake_case_       (			self,			SCREAMING_SNAKE_CASE_,			SCREAMING_SNAKE_CASE_,			SCREAMING_SNAKE_CASE_       )     ->		Tuple:
          UpperCamelCase    :  Dict							=	ConvNextModel(config=SCREAMING_SNAKE_CASE_       )
          model.to(SCREAMING_SNAKE_CASE_       )
          model.eval()
          UpperCamelCase    :  str							=	model(SCREAMING_SNAKE_CASE_       )
          # 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,			SCREAMING_SNAKE_CASE_,			SCREAMING_SNAKE_CASE_,			SCREAMING_SNAKE_CASE_       )     ->		Optional[Any]:
          UpperCamelCase    :  Optional[int]							=	ConvNextForImageClassification(SCREAMING_SNAKE_CASE_       )
          model.to(SCREAMING_SNAKE_CASE_       )
          model.eval()
          UpperCamelCase    :  Dict							=	model(SCREAMING_SNAKE_CASE_,			labels=SCREAMING_SNAKE_CASE_       )
          self.parent.assertEqual(result.logits.shape,			(self.batch_size, self.num_labels)       )
       def 					snake_case_       (			self,			SCREAMING_SNAKE_CASE_,			SCREAMING_SNAKE_CASE_,			SCREAMING_SNAKE_CASE_       )     ->		Dict:
          UpperCamelCase    :  int							=	ConvNextBackbone(config=SCREAMING_SNAKE_CASE_       )
          model.to(SCREAMING_SNAKE_CASE_       )
          model.eval()
          UpperCamelCase    :  Any							=	model(SCREAMING_SNAKE_CASE_       )

          # verify hidden states
          self.parent.assertEqual(len(result.feature_maps       ),			len(config.out_features       )       )
          self.parent.assertListEqual(list(result.feature_maps[0].shape       ),			[self.batch_size, self.hidden_sizes[1], 4, 4]       )

          # verify channels
          self.parent.assertEqual(len(model.channels       ),			len(config.out_features       )       )
          self.parent.assertListEqual(model.channels,			config.hidden_sizes[1:]       )

          # verify backbone works with out_features=None
          UpperCamelCase    :  Any							=	None
          UpperCamelCase    :  List[str]							=	ConvNextBackbone(config=SCREAMING_SNAKE_CASE_       )
          model.to(SCREAMING_SNAKE_CASE_       )
          model.eval()
          UpperCamelCase    :  Tuple							=	model(SCREAMING_SNAKE_CASE_       )

          # verify feature maps
          self.parent.assertEqual(len(result.feature_maps       ),			1       )
          self.parent.assertListEqual(list(result.feature_maps[0].shape       ),			[self.batch_size, self.hidden_sizes[-1], 1, 1]       )

          # verify channels
          self.parent.assertEqual(len(model.channels       ),			1       )
          self.parent.assertListEqual(model.channels,			[config.hidden_sizes[-1]]       )






       def 					snake_case_       (			self       )     ->		int:
          UpperCamelCase    :  Optional[int]							=	self.prepare_config_and_inputs()
          UpperCamelCase	,     UpperCamelCase	,     UpperCamelCase    :  Optional[Any]							=	config_and_inputs
          UpperCamelCase    :  str							=	{'pixel_values': pixel_values}
          return config, inputs_dict






@require_torch
class 				lowerCAmelCase_		(  a__       ,    a__       ,    unittest.TestCase							):
       UpperCAmelCase__				:					List[Any]        =  (
           (
               ConvNextModel,
               ConvNextForImageClassification,
               ConvNextBackbone,
           )
           if is_torch_available()
           else ()
       )
       UpperCAmelCase__				:					Union[str, Any]        =  (
           {"feature-extraction": ConvNextModel, "image-classification": ConvNextForImageClassification}
           if is_torch_available()
           else {}
       )

       UpperCAmelCase__				:					Dict        =  True
       UpperCAmelCase__				:					List[str]        =  False
       UpperCAmelCase__				:					Tuple        =  False
       UpperCAmelCase__				:					List[Any]        =  False
       UpperCAmelCase__				:					Tuple        =  False
       def 					snake_case_       (			self       )     ->		Optional[Any]:
          UpperCamelCase    :  Dict							=	ConvNextModelTester(self       )
          UpperCamelCase    :  int							=	ConfigTester(self,			config_class=SCREAMING_SNAKE_CASE_,			has_text_modality=SCREAMING_SNAKE_CASE_,			hidden_size=37       )
       def 					snake_case_       (			self       )     ->		List[str]:
          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       )     ->		Union[str, Any]:
          return
       @unittest.skip(reason='ConvNext does not use inputs_embeds'       )
       def 					snake_case_       (			self       )     ->		List[str]:
          pass
       @unittest.skip(reason='ConvNext does not support input and output embeddings'       )
       def 					snake_case_       (			self       )     ->		str:
          pass
       @unittest.skip(reason='ConvNext does not use feedforward chunking'       )
       def 					snake_case_       (			self       )     ->		str:
          pass
       def 					snake_case_       (			self       )     ->		Any:
          UpperCamelCase	,     UpperCamelCase    :  Optional[int]							=	self.model_tester.prepare_config_and_inputs_for_common()

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

             UpperCamelCase    :  Union[str, Any]							=	['pixel_values']
             self.assertListEqual(arg_names[:1],			SCREAMING_SNAKE_CASE_       )
       def 					snake_case_       (			self       )     ->		List[str]:
          UpperCamelCase    :  Optional[int]							=	self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_model(*SCREAMING_SNAKE_CASE_       )
       def 					snake_case_       (			self       )     ->		Dict:
          UpperCamelCase    :  Optional[Any]							=	self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_backbone(*SCREAMING_SNAKE_CASE_       )
       def 					snake_case_       (			self       )     ->		Dict:
          def check_hidden_states_output(SCREAMING_SNAKE_CASE_,			SCREAMING_SNAKE_CASE_,			SCREAMING_SNAKE_CASE_       ):
             UpperCamelCase    :  List[Any]							=	model_class(SCREAMING_SNAKE_CASE_       )
             model.to(SCREAMING_SNAKE_CASE_       )
             model.eval()

             with torch.no_grad():
                UpperCamelCase    :  List[Any]							=	model(**self._prepare_for_class(SCREAMING_SNAKE_CASE_,			SCREAMING_SNAKE_CASE_       )       )

             UpperCamelCase    :  List[str]							=	outputs.encoder_hidden_states if config.is_encoder_decoder else outputs.hidden_states

             UpperCamelCase    :  int							=	self.model_tester.num_stages
             self.assertEqual(len(SCREAMING_SNAKE_CASE_       ),			expected_num_stages + 1       )

             # ConvNext'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],			)

          UpperCamelCase	,     UpperCamelCase    :  Tuple							=	self.model_tester.prepare_config_and_inputs_for_common()

          for model_class in self.all_model_classes:
             UpperCamelCase    :  Optional[int]							=	True
             check_hidden_states_output(SCREAMING_SNAKE_CASE_,			SCREAMING_SNAKE_CASE_,			SCREAMING_SNAKE_CASE_       )

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

             check_hidden_states_output(SCREAMING_SNAKE_CASE_,			SCREAMING_SNAKE_CASE_,			SCREAMING_SNAKE_CASE_       )
       def 					snake_case_       (			self       )     ->		List[str]:
          UpperCamelCase    :  List[Any]							=	self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_for_image_classification(*SCREAMING_SNAKE_CASE_       )






       @slow
       def 					snake_case_       (			self       )     ->		Tuple:
          for model_name in CONVNEXT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
             UpperCamelCase    :  Union[str, Any]							=	ConvNextModel.from_pretrained(SCREAMING_SNAKE_CASE_       )
             self.assertIsNotNone(SCREAMING_SNAKE_CASE_       )






def 		UpperCamelCase	(   )	->     int:
   UpperCamelCase    :  Optional[int]							=	Image.open('./tests/fixtures/tests_samples/COCO/000000039769.png'	)
   return image



@require_torch
@require_vision
class 				lowerCAmelCase_		(  unittest.TestCase							):
       @cached_property
       def 					snake_case_       (			self       )     ->		Optional[int]:
          return AutoImageProcessor.from_pretrained('facebook/convnext-tiny-224'       ) if is_vision_available() else None






       @slow
       def 					snake_case_       (			self       )     ->		str:
          UpperCamelCase    :  Dict							=	ConvNextForImageClassification.from_pretrained('facebook/convnext-tiny-224'       ).to(SCREAMING_SNAKE_CASE_       )

          UpperCamelCase    :  Union[str, Any]							=	self.default_image_processor
          UpperCamelCase    :  Any							=	prepare_img()
          UpperCamelCase    :  List[str]							=	image_processor(images=SCREAMING_SNAKE_CASE_,			return_tensors='pt'       ).to(SCREAMING_SNAKE_CASE_       )

          # forward pass
          with torch.no_grad():
             UpperCamelCase    :  Optional[Any]							=	model(**SCREAMING_SNAKE_CASE_       )

          # verify the logits
          UpperCamelCase    :  List[Any]							=	torch.Size((1, 1000)       )
          self.assertEqual(outputs.logits.shape,			SCREAMING_SNAKE_CASE_       )

          UpperCamelCase    :  List[str]							=	torch.tensor([-0.02_60, -0.47_39, 0.19_11]       ).to(SCREAMING_SNAKE_CASE_       )

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






@require_torch
class 				lowerCAmelCase_		(  unittest.TestCase       ,    a__							):
       UpperCAmelCase__				:					Tuple        =  (ConvNextBackbone,) if is_torch_available() else ()
       UpperCAmelCase__				:					List[str]        =  ConvNextConfig

       UpperCAmelCase__				:					Tuple        =  False
       def 					snake_case_       (			self       )     ->		int:
          UpperCamelCase    :  List[Any]							=	ConvNextModelTester(self       )



