

import inspect
import unittest

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

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


if is_torch_available():
 import torch
 from torch import nn

 from transformers import SwinaSRForImageSuperResolution, SwinaSRModel
 from transformers.models.swinasr.modeling_swinasr import SWIN2SR_PRETRAINED_MODEL_ARCHIVE_LIST

if is_vision_available():
 from PIL import Image

 from transformers import SwinaSRImageProcessor






class 		__magic_name__ :






     def __init__(							self						,		_lowercase						,		_lowercase=13						,		_lowercase=32						,		_lowercase=1						,		_lowercase=3						,		_lowercase=16						,		_lowercase=[1, 2, 1]						,		_lowercase=[2, 2, 4]						,		_lowercase=2						,		_lowercase=2.0						,		_lowercase=True						,		_lowercase=0.0						,		_lowercase=0.0						,		_lowercase=0.1						,		_lowercase="gelu"						,		_lowercase=False						,		_lowercase=True						,		_lowercase=0.02						,		_lowercase=1e-5						,		_lowercase=True						,		_lowercase=None						,		_lowercase=False						,		_lowercase=2						,		)->			List[Any]:
         UpperCamelCase_          =		parent
         UpperCamelCase_          =		batch_size
         UpperCamelCase_          =		image_size
         UpperCamelCase_          =		patch_size
         UpperCamelCase_          =		num_channels
         UpperCamelCase_          =		embed_dim
         UpperCamelCase_          =		depths
         UpperCamelCase_          =		num_heads
         UpperCamelCase_          =		window_size
         UpperCamelCase_          =		mlp_ratio
         UpperCamelCase_          =		qkv_bias
         UpperCamelCase_          =		hidden_dropout_prob
         UpperCamelCase_          =		attention_probs_dropout_prob
         UpperCamelCase_          =		drop_path_rate
         UpperCamelCase_          =		hidden_act
         UpperCamelCase_          =		use_absolute_embeddings
         UpperCamelCase_          =		patch_norm
         UpperCamelCase_          =		layer_norm_eps
         UpperCamelCase_          =		initializer_range
         UpperCamelCase_          =		is_training
         UpperCamelCase_          =		scope
         UpperCamelCase_          =		use_labels
         UpperCamelCase_          =		upscale

         # here we set some attributes to make tests pass
         UpperCamelCase_          =		len(_lowercase     )
         UpperCamelCase_          =		embed_dim
         UpperCamelCase_          =		(image_size // patch_size) ** 2






     def      UpperCAmelCase_							(							self     )->			Any:
         UpperCamelCase_          =		floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]     )

         UpperCamelCase_          =		None
         if self.use_labels:
             UpperCamelCase_          =		ids_tensor([self.batch_size]						,		self.type_sequence_label_size     )

         UpperCamelCase_          =		self.get_config()

         return config, pixel_values, labels






     def      UpperCAmelCase_							(							self     )->			str:
         return SwinaSRConfig(
             image_size=self.image_size						,		patch_size=self.patch_size						,		num_channels=self.num_channels						,		embed_dim=self.embed_dim						,		depths=self.depths						,		num_heads=self.num_heads						,		window_size=self.window_size						,		mlp_ratio=self.mlp_ratio						,		qkv_bias=self.qkv_bias						,		hidden_dropout_prob=self.hidden_dropout_prob						,		attention_probs_dropout_prob=self.attention_probs_dropout_prob						,		drop_path_rate=self.drop_path_rate						,		hidden_act=self.hidden_act						,		use_absolute_embeddings=self.use_absolute_embeddings						,		path_norm=self.patch_norm						,		layer_norm_eps=self.layer_norm_eps						,		initializer_range=self.initializer_range						,		upscale=self.upscale						,		)






     def      UpperCAmelCase_							(							self						,		_lowercase						,		_lowercase						,		_lowercase     )->			Union[str, Any]:
         UpperCamelCase_          =		SwinaSRModel(config=_lowercase     )
         model.to(_lowercase     )
         model.eval()
         UpperCamelCase_          =		model(_lowercase     )

         self.parent.assertEqual(
             result.last_hidden_state.shape						,		(self.batch_size, self.embed_dim, self.image_size, self.image_size)     )






     def      UpperCAmelCase_							(							self						,		_lowercase						,		_lowercase						,		_lowercase     )->			Optional[Any]:
         UpperCamelCase_          =		SwinaSRForImageSuperResolution(_lowercase     )
         model.to(_lowercase     )
         model.eval()
         UpperCamelCase_          =		model(_lowercase     )

         UpperCamelCase_          =		self.image_size * self.upscale
         self.parent.assertEqual(
             result.reconstruction.shape						,		(self.batch_size, self.num_channels, expected_image_size, expected_image_size)     )






     def      UpperCAmelCase_							(							self     )->			Tuple:
         UpperCamelCase_          =		self.prepare_config_and_inputs()
         UpperCamelCase_   ,    UpperCamelCase_   ,    UpperCamelCase_          =		config_and_inputs
         UpperCamelCase_          =		{"pixel_values": pixel_values}
         return config, inputs_dict






@require_torch
class 		__magic_name__ (				snake_case       ,	snake_case       ,	unittest.TestCase							):
     UpperCamelCase_     :Any									=	(SwinaSRModel, SwinaSRForImageSuperResolution) if is_torch_available() else ()
     UpperCamelCase_     :Optional[Any]									=	{"""feature-extraction""": SwinaSRModel} if is_torch_available() else {}

     UpperCamelCase_     :List[Any]									=	False
     UpperCamelCase_     :int									=	False
     UpperCamelCase_     :Optional[Any]									=	False
     UpperCamelCase_     :Optional[Any]									=	False
     UpperCamelCase_     :Tuple									=	False






     def      UpperCAmelCase_							(							self     )->			Any:
         UpperCamelCase_          =		SwinaSRModelTester(self     )
         UpperCamelCase_          =		ConfigTester(self						,		config_class=_lowercase						,		embed_dim=37     )






     def      UpperCAmelCase_							(							self     )->			str:
         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      UpperCAmelCase_							(							self     )->			List[Any]:
         UpperCamelCase_          =		self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_model(*_lowercase     )






     def      UpperCAmelCase_							(							self     )->			Any:
         UpperCamelCase_          =		self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_image_super_resolution(*_lowercase     )






     @unittest.skip(reason="Got `CUDA error: misaligned address` with PyTorch 2.0.0."     )
     def      UpperCAmelCase_							(							self     )->			Optional[int]:
         pass






     @unittest.skip(reason="Swin2SR does not use inputs_embeds"     )
     def      UpperCAmelCase_							(							self     )->			int:
         pass






     @unittest.skip(reason="Swin2SR does not support training yet"     )
     def      UpperCAmelCase_							(							self     )->			Union[str, Any]:
         pass






     @unittest.skip(reason="Swin2SR does not support training yet"     )
     def      UpperCAmelCase_							(							self     )->			Dict:
         pass






     def      UpperCAmelCase_							(							self     )->			Union[str, Any]:
         UpperCamelCase_   ,    UpperCamelCase_          =		self.model_tester.prepare_config_and_inputs_for_common()

         for model_class in self.all_model_classes:
             UpperCamelCase_          =		model_class(_lowercase     )
             self.assertIsInstance(model.get_input_embeddings()						,		(nn.Module)     )
             UpperCamelCase_          =		model.get_output_embeddings()
             self.assertTrue(x is None or isinstance(_lowercase						,		nn.Linear     )     )






     def      UpperCAmelCase_							(							self     )->			Optional[int]:
         UpperCamelCase_   ,    UpperCamelCase_          =		self.model_tester.prepare_config_and_inputs_for_common()

         for model_class in self.all_model_classes:
             UpperCamelCase_          =		model_class(_lowercase     )
             UpperCamelCase_          =		inspect.signature(model.forward     )
             # signature.parameters is an OrderedDict => so arg_names order is deterministic
             UpperCamelCase_          =		[*signature.parameters.keys()]

             UpperCamelCase_          =		["pixel_values"]
             self.assertListEqual(arg_names[:1]						,		_lowercase     )






     @slow
     def      UpperCAmelCase_							(							self     )->			Optional[int]:
         for model_name in SWIN2SR_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
             UpperCamelCase_          =		SwinaSRModel.from_pretrained(_lowercase     )
             self.assertIsNotNone(_lowercase     )






     def      UpperCAmelCase_							(							self     )->			Optional[Any]:
         UpperCamelCase_   ,    UpperCamelCase_          =		self.model_tester.prepare_config_and_inputs_for_common()

         UpperCamelCase_          =		_config_zero_init(_lowercase     )
         for model_class in self.all_model_classes:
             UpperCamelCase_          =		model_class(config=_lowercase     )
             for name, param in model.named_parameters():
                 if "logit_scale" in name:
                     continue
                 if param.requires_grad:
                     self.assertIn(
                         ((param.data.mean() * 1e9).round() / 1e9).item()						,		[0.0, 1.0]						,		msg=F"Parameter {name} of model {model_class} seems not properly initialized"						,		)






     def      UpperCAmelCase_							(							self     )->			Any:
         UpperCamelCase_   ,    UpperCamelCase_          =		self.model_tester.prepare_config_and_inputs_for_common()
         UpperCamelCase_          =		True

         for model_class in self.all_model_classes:
             UpperCamelCase_          =		True
             UpperCamelCase_          =		False
             UpperCamelCase_          =		True
             UpperCamelCase_          =		model_class(_lowercase     )
             model.to(_lowercase     )
             model.eval()
             with torch.no_grad():
                 UpperCamelCase_          =		model(**self._prepare_for_class(_lowercase						,		_lowercase     )     )
             UpperCamelCase_          =		outputs.attentions
             UpperCamelCase_          =		len(self.model_tester.depths     )
             self.assertEqual(len(_lowercase     )						,		_lowercase     )

             # check that output_attentions also work using config
             del inputs_dict["output_attentions"]
             UpperCamelCase_          =		True
             UpperCamelCase_          =		config.window_size**2
             UpperCamelCase_          =		model_class(_lowercase     )
             model.to(_lowercase     )
             model.eval()
             with torch.no_grad():
                 UpperCamelCase_          =		model(**self._prepare_for_class(_lowercase						,		_lowercase     )     )
             UpperCamelCase_          =		outputs.attentions
             self.assertEqual(len(_lowercase     )						,		_lowercase     )

             self.assertListEqual(
                 list(attentions[0].shape[-3:]     )						,		[self.model_tester.num_heads[0], window_size_squared, window_size_squared]						,		)
             UpperCamelCase_          =		len(_lowercase     )

             # Check attention is always last and order is fine
             UpperCamelCase_          =		True
             UpperCamelCase_          =		True
             UpperCamelCase_          =		model_class(_lowercase     )
             model.to(_lowercase     )
             model.eval()
             with torch.no_grad():
                 UpperCamelCase_          =		model(**self._prepare_for_class(_lowercase						,		_lowercase     )     )

             self.assertEqual(out_len + 1						,		len(_lowercase     )     )

             UpperCamelCase_          =		outputs.attentions

             self.assertEqual(len(_lowercase     )						,		_lowercase     )

             self.assertListEqual(
                 list(self_attentions[0].shape[-3:]     )						,		[self.model_tester.num_heads[0], window_size_squared, window_size_squared]						,		)






@require_vision
@require_torch
@slow
class 		__magic_name__ (				unittest.TestCase							):






     def      UpperCAmelCase_							(							self     )->			Optional[int]:
         UpperCamelCase_          =		SwinaSRImageProcessor()
         UpperCamelCase_          =		SwinaSRForImageSuperResolution.from_pretrained("caidas/swin2SR-classical-sr-x2-64"     ).to(_lowercase     )

         UpperCamelCase_          =		Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png"     )
         UpperCamelCase_          =		processor(images=_lowercase						,		return_tensors="pt"     ).to(_lowercase     )

         # forward pass
         with torch.no_grad():
             UpperCamelCase_          =		model(**_lowercase     )

         # verify the logits
         UpperCamelCase_          =		torch.Size([1, 3, 976, 1_296]     )
         self.assertEqual(outputs.reconstruction.shape						,		_lowercase     )
         UpperCamelCase_          =		torch.tensor(
             [[0.5_458, 0.5_546, 0.5_638], [0.5_526, 0.5_565, 0.5_651], [0.5_396, 0.5_426, 0.5_621]]     ).to(_lowercase     )
         self.assertTrue(torch.allclose(outputs.reconstruction[0, 0, :3, :3]						,		_lowercase						,		atol=1e-4     )     )





"""simple docstring"""







import numpy as np
import torch
import torch.nn as nn
from transformers import CLIPConfig, CLIPVisionModel, PreTrainedModel

from ...utils import logging


UpperCAmelCase:   Dict       			=							logging.get_logger(__name__)


def       __SCREAMING_SNAKE_CASE      (   __UpperCAmelCase					,							__UpperCAmelCase	):
		_lowercase							:      List[str]			=					nn.functional.normalize(__UpperCAmelCase	)
		_lowercase							:      int			=					nn.functional.normalize(__UpperCAmelCase	)
		return torch.mm(__UpperCAmelCase					,							normalized_text_embeds.t()	)





class    UpperCamelCase     (			snake_case    ):




	"""simple docstring"""
	SCREAMING_SNAKE_CASE_				:    List[str]		    =   CLIPConfig

	SCREAMING_SNAKE_CASE_				:    List[Any]		    =   ["CLIPEncoderLayer"]






	def __init__(			self							,UpperCAmelCase_					):
			super().__init__(UpperCAmelCase_					)

			_lowercase							:      Any			=					CLIPVisionModel(config.vision_config					)
			_lowercase							:      Any			=					nn.Linear(config.vision_config.hidden_size							,config.projection_dim							,bias=UpperCAmelCase_					)

			_lowercase							:      Dict			=					nn.Parameter(torch.ones(17							,config.projection_dim					)							,requires_grad=UpperCAmelCase_					)
			_lowercase							:      Tuple			=					nn.Parameter(torch.ones(3							,config.projection_dim					)							,requires_grad=UpperCAmelCase_					)

			_lowercase							:      Optional[int]			=					nn.Parameter(torch.ones(17					)							,requires_grad=UpperCAmelCase_					)
			_lowercase							:      List[Any]			=					nn.Parameter(torch.ones(3					)							,requires_grad=UpperCAmelCase_					)






	@torch.no_grad()
	def 					lowerCamelCase__ (			self							,UpperCAmelCase_							,UpperCAmelCase_					):
			_lowercase							:      List[str]			=					self.vision_model(UpperCAmelCase_					)[1]  # pooled_output
			_lowercase							:      str			=					self.visual_projection(UpperCAmelCase_					)

			# we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16
			_lowercase							:      Union[str, Any]			=					cosine_distance(UpperCAmelCase_							,self.special_care_embeds					).cpu().float().numpy()
			_lowercase							:      str			=					cosine_distance(UpperCAmelCase_							,self.concept_embeds					).cpu().float().numpy()

			_lowercase							:      List[str]			=					[]
			_lowercase							:      Optional[Any]			=					image_embeds.shape[0]
			for i in range(UpperCAmelCase_					):
					_lowercase							:      int			=					{"""special_scores""": {}, """special_care""": [], """concept_scores""": {}, """bad_concepts""": []}

					# increase this value to create a stronger `nfsw` filter
					# at the cost of increasing the possibility of filtering benign images
					_lowercase							:      str			=					0.0

					for concept_idx in range(len(special_cos_dist[0]					)					):
							_lowercase							:      Optional[Any]			=					special_cos_dist[i][concept_idx]
							_lowercase							:      str			=					self.special_care_embeds_weights[concept_idx].item()
							_lowercase							:      Tuple			=					round(concept_cos - concept_threshold + adjustment							,3					)
							if result_img["special_scores"][concept_idx] > 0:
									result_img["special_care"].append({concept_idx, result_img["""special_scores"""][concept_idx]}					)
									_lowercase							:      Union[str, Any]			=					0.01

					for concept_idx in range(len(cos_dist[0]					)					):
							_lowercase							:      Union[str, Any]			=					cos_dist[i][concept_idx]
							_lowercase							:      str			=					self.concept_embeds_weights[concept_idx].item()
							_lowercase							:      str			=					round(concept_cos - concept_threshold + adjustment							,3					)
							if result_img["concept_scores"][concept_idx] > 0:
									result_img["bad_concepts"].append(UpperCAmelCase_					)

					result.append(UpperCAmelCase_					)

			_lowercase							:      List[str]			=					[len(res["""bad_concepts"""]					) > 0 for res in result]

			for idx, has_nsfw_concept in enumerate(UpperCAmelCase_					):
					if has_nsfw_concept:
							if torch.is_tensor(UpperCAmelCase_					) or torch.is_tensor(images[0]					):
									_lowercase							:      List[Any]			=					torch.zeros_like(images[idx]					)  # black image
							else:
									_lowercase							:      Any			=					np.zeros(images[idx].shape					)  # black image

			if any(UpperCAmelCase_					):
					logger.warning(
					    """Potential NSFW content was detected in one or more images. A black image will be returned instead."""
					    """ Try again with a different prompt and/or seed."""					)

			return images, has_nsfw_concepts






	@torch.no_grad()
	def 					lowerCamelCase__ (			self							,UpperCAmelCase_							,UpperCAmelCase_					):
			_lowercase							:      List[Any]			=					self.vision_model(UpperCAmelCase_					)[1]  # pooled_output
			_lowercase							:      Optional[int]			=					self.visual_projection(UpperCAmelCase_					)

			_lowercase							:      Union[str, Any]			=					cosine_distance(UpperCAmelCase_							,self.special_care_embeds					)
			_lowercase							:      Tuple			=					cosine_distance(UpperCAmelCase_							,self.concept_embeds					)

			# increase this value to create a stronger `nsfw` filter
			# at the cost of increasing the possibility of filtering benign images
			_lowercase							:      Any			=					0.0

			_lowercase							:      Tuple			=					special_cos_dist - self.special_care_embeds_weights + adjustment
			# special_scores = special_scores.round(decimals=3)
			_lowercase							:      Dict			=					torch.any(special_scores > 0							,dim=1					)
			_lowercase							:      Any			=					special_care * 0.01
			_lowercase							:      Optional[int]			=					special_adjustment.unsqueeze(1					).expand(-1							,cos_dist.shape[1]					)

			_lowercase							:      Optional[Any]			=					(cos_dist - self.concept_embeds_weights) + special_adjustment
			# concept_scores = concept_scores.round(decimals=3)
			_lowercase							:      Optional[int]			=					torch.any(concept_scores > 0							,dim=1					)

			_lowercase							:      Tuple			=					0.0  # black image

			return images, has_nsfw_concepts



