

import warnings

from diffusers import StableDiffusionImgaImgPipeline  # noqa F401


warnings.warn(
    "The `image_to_image.py` script is outdated. Please use directly `from diffusers import"
    " StableDiffusionImg2ImgPipeline` instead."
)








import inspect
import unittest

from transformers import BitConfig
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 torch import nn

     from transformers import BitBackbone, BitForImageClassification, BitImageProcessor, BitModel
     from transformers.models.bit.modeling_bit import BIT_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
     from PIL import Image





class  __UpperCAmelCase					:



   def __init__(							self							:					Any,		__A							:					Optional[Any],		__A							:					Optional[Any]=3,		__A							:					Union[str, Any]=3_2,		__A							:					Optional[int]=3,		__A							:					str=1_0,		__A							:					Union[str, Any]=[8, 1_6, 3_2, 6_4],		__A							:					List[str]=[1, 1, 2, 1],		__A							:					Dict=True,		__A							:					List[Any]=True,		__A							:					int="relu",		__A							:					Optional[Any]=3,		__A							:					Any=None,		__A							:					Any=["stage2", "stage3", "stage4"],		__A							:					Optional[int]=[2, 3, 4],		__A							:					Any=1,		):
     UpperCAmelCase	:							Optional[int]						=			parent
     UpperCAmelCase	:							Union[str, Any]						=			batch_size
     UpperCAmelCase	:							Optional[int]						=			image_size
     UpperCAmelCase	:							List[Any]						=			num_channels
     UpperCAmelCase	:							Optional[Any]						=			embeddings_size
     UpperCAmelCase	:							List[Any]						=			hidden_sizes
     UpperCAmelCase	:							Any						=			depths
     UpperCAmelCase	:							int						=			is_training
     UpperCAmelCase	:							Optional[int]						=			use_labels
     UpperCAmelCase	:							Optional[Any]						=			hidden_act
     UpperCAmelCase	:							Any						=			num_labels
     UpperCAmelCase	:							List[Any]						=			scope
     UpperCAmelCase	:							int						=			len(__A						)
     UpperCAmelCase	:							Union[str, Any]						=			out_features
     UpperCAmelCase	:							List[str]						=			out_indices
     UpperCAmelCase	:							Tuple						=			num_groups



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

     UpperCAmelCase	:							Dict						=			None
     if self.use_labels:
       UpperCAmelCase	:							List[str]						=			ids_tensor([self.batch_size],		self.num_labels						)

     UpperCAmelCase	:							List[str]						=			self.get_config()

     return config, pixel_values, labels



   def 	__magic_name__   (							self							:					Tuple						):
     return BitConfig(
         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,		out_features=self.out_features,		out_indices=self.out_indices,		num_groups=self.num_groups,		)



   def 	__magic_name__   (							self							:					int,		__A							:					str,		__A							:					List[Any],		__A							:					Any						):
     UpperCAmelCase	:							Optional[int]						=			BitModel(config=__A						)
     model.to(__A						)
     model.eval()
     UpperCAmelCase	:							Optional[int]						=			model(__A						)
     self.parent.assertEqual(
         result.last_hidden_state.shape,		(self.batch_size, self.hidden_sizes[-1], self.image_size // 3_2, self.image_size // 3_2),		)



   def 	__magic_name__   (							self							:					List[Any],		__A							:					Any,		__A							:					Union[str, Any],		__A							:					Dict						):
     UpperCAmelCase	:							Optional[Any]						=			self.num_labels
     UpperCAmelCase	:							Union[str, Any]						=			BitForImageClassification(__A						)
     model.to(__A						)
     model.eval()
     UpperCAmelCase	:							Optional[int]						=			model(__A,		labels=__A						)
     self.parent.assertEqual(result.logits.shape,		(self.batch_size, self.num_labels)						)



   def 	__magic_name__   (							self							:					Optional[int],		__A							:					Any,		__A							:					List[str],		__A							:					str						):
     UpperCAmelCase	:							Any						=			BitBackbone(config=__A						)
     model.to(__A						)
     model.eval()
     UpperCAmelCase	:							Tuple						=			model(__A						)

     # verify feature maps
     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	:							List[str]						=			None
     UpperCAmelCase	:							Optional[Any]						=			BitBackbone(config=__A						)
     model.to(__A						)
     model.eval()
     UpperCAmelCase	:							Tuple						=			model(__A						)

     # 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 	__magic_name__   (							self							:					Optional[int]						):
     UpperCAmelCase	:							str						=			self.prepare_config_and_inputs()
     UpperCAmelCase			,						UpperCAmelCase			,						UpperCAmelCase	:							Union[str, Any]						=			config_and_inputs
     UpperCAmelCase	:							Union[str, Any]						=			{'''pixel_values''': pixel_values}
     return config, inputs_dict






@require_torch
class  __UpperCAmelCase					(	lowerCamelCase__      ,  lowerCamelCase__      ,  unittest.TestCase						):
   UpperCamelCase       					=					(BitModel, BitForImageClassification, BitBackbone) if is_torch_available() else ()
   UpperCamelCase       					=					(
       {"""feature-extraction""": BitModel, """image-classification""": BitForImageClassification}
       if is_torch_available()
       else {}
   )

   UpperCamelCase       					=					False
   UpperCamelCase       					=					False
   UpperCamelCase       					=					False
   UpperCamelCase       					=					False
   UpperCamelCase       					=					False



   def 	__magic_name__   (							self							:					Union[str, Any]						):
     UpperCAmelCase	:							Optional[int]						=			BitModelTester(self						)
     UpperCAmelCase	:							Any						=			ConfigTester(self,		config_class=__A,		has_text_modality=__A						)



   def 	__magic_name__   (							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 	__magic_name__   (							self							:					Union[str, Any]						):
     return



   @unittest.skip(reason='''Bit does not output attentions'''						)
   def 	__magic_name__   (							self							:					List[Any]						):
     pass



   @unittest.skip(reason='''Bit does not use inputs_embeds'''						)
   def 	__magic_name__   (							self							:					Dict						):
     pass



   @unittest.skip(reason='''Bit does not support input and output embeddings'''						)
   def 	__magic_name__   (							self							:					List[Any]						):
     pass



   def 	__magic_name__   (							self							:					List[Any]						):
     UpperCAmelCase			,						UpperCAmelCase	:							Any						=			self.model_tester.prepare_config_and_inputs_for_common()

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

       UpperCAmelCase	:							List[Any]						=			['''pixel_values''']
       self.assertListEqual(arg_names[:1],		__A						)



   def 	__magic_name__   (							self							:					Union[str, Any]						):
     UpperCAmelCase	:							Optional[int]						=			self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_model(*__A						)



   def 	__magic_name__   (							self							:					int						):
     UpperCAmelCase	:							List[str]						=			self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_backbone(*__A						)



   def 	__magic_name__   (							self							:					int						):
     UpperCAmelCase			,						UpperCAmelCase	:							Optional[int]						=			self.model_tester.prepare_config_and_inputs_for_common()

     for model_class in self.all_model_classes:
       UpperCAmelCase	:							Dict						=			model_class(config=__A						)
       for name, module in model.named_modules():
         if isinstance(__A,		(nn.BatchNormad, nn.GroupNorm)						):
           self.assertTrue(
               torch.all(module.weight == 1						),		msg=F'''Parameter {name} of model {model_class} seems not properly initialized''',		)
           self.assertTrue(
               torch.all(module.bias == 0						),		msg=F'''Parameter {name} of model {model_class} seems not properly initialized''',		)



   def 	__magic_name__   (							self							:					Dict						):
     def check_hidden_states_output(__A							:					List[Any],		__A							:					Optional[int],		__A							:					int						):
       UpperCAmelCase	:							int						=			model_class(__A						)
       model.to(__A						)
       model.eval()

       with torch.no_grad():
         UpperCAmelCase	:							Any						=			model(**self._prepare_for_class(__A,		__A						)						)

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

       UpperCAmelCase	:							Any						=			self.model_tester.num_stages
       self.assertEqual(len(__A						),		expected_num_stages + 1						)

       # Bit'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	:							Dict						=			self.model_tester.prepare_config_and_inputs_for_common()
     UpperCAmelCase	:							Union[str, Any]						=			['''preactivation''', '''bottleneck''']
     for model_class in self.all_model_classes:
       for layer_type in layers_type:
         UpperCAmelCase	:							int						=			layer_type
         UpperCAmelCase	:							List[Any]						=			True
         check_hidden_states_output(__A,		__A,		__A						)

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

         check_hidden_states_output(__A,		__A,		__A						)



   @unittest.skip(reason='''Bit does not use feedforward chunking'''						)
   def 	__magic_name__   (							self							:					List[str]						):
     pass



   def 	__magic_name__   (							self							:					Optional[int]						):
     UpperCAmelCase	:							str						=			self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_image_classification(*__A						)





   @slow
   def 	__magic_name__   (							self							:					List[Any]						):
     for model_name in BIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
       UpperCAmelCase	:							List[str]						=			BitModel.from_pretrained(__A						)
       self.assertIsNotNone(__A						)






def       a__					(		)  ->						Union[str, Any]:
  UpperCAmelCase	:							Optional[int]						=			Image.open('''./tests/fixtures/tests_samples/COCO/000000039769.png'''      )
  return image





@require_torch
@require_vision
class  __UpperCAmelCase					(	unittest.TestCase						):



   @cached_property
   def 	__magic_name__   (							self							:					Union[str, Any]						):
     return (
         BitImageProcessor.from_pretrained(BIT_PRETRAINED_MODEL_ARCHIVE_LIST[0]						) if is_vision_available() else None
     )





   @slow
   def 	__magic_name__   (							self							:					str						):
     UpperCAmelCase	:							List[Any]						=			BitForImageClassification.from_pretrained(BIT_PRETRAINED_MODEL_ARCHIVE_LIST[0]						).to(__A						)

     UpperCAmelCase	:							str						=			self.default_image_processor
     UpperCAmelCase	:							Dict						=			prepare_img()
     UpperCAmelCase	:							Tuple						=			image_processor(images=__A,		return_tensors='''pt'''						).to(__A						)

     # forward pass
     with torch.no_grad():
       UpperCAmelCase	:							List[Any]						=			model(**__A						)

     # verify the logits
     UpperCAmelCase	:							Dict						=			torch.Size((1, 1_0_0_0)						)
     self.assertEqual(outputs.logits.shape,		__A						)

     UpperCAmelCase	:							int						=			torch.tensor([[-0.6_5_2_6, -0.5_2_6_3, -1.4_3_9_8]]						).to(__A						)

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






@require_torch
class  __UpperCAmelCase					(	lowerCamelCase__      ,  unittest.TestCase						):
   UpperCamelCase       					=					(BitBackbone,) if is_torch_available() else ()
   UpperCamelCase       					=					BitConfig

   UpperCamelCase       					=					False



   def 	__magic_name__   (							self							:					int						):
     UpperCAmelCase	:							int						=			BitModelTester(self						)





