






import inspect
import unittest

from transformers import ViTHybridConfig
from transformers.testing_utils import require_accelerate, require_torch, require_vision, slow, torch_device
from transformers.utils import cached_property, 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 ViTHybridForImageClassification, ViTHybridImageProcessor, ViTHybridModel
							from transformers.models.vit_hybrid.modeling_vit_hybrid import VIT_HYBRID_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
							from PIL import Image


class   __SCREAMING_SNAKE_CASE:
							def __init__(     self:      Tuple       ,	UpperCamelCase:      Optional[int]       ,	UpperCamelCase:      Dict=13       ,	UpperCamelCase:      Union[str, Any]=64       ,	UpperCamelCase:      Optional[Any]=2       ,	UpperCamelCase:      str=3       ,	UpperCamelCase:      Optional[int]=True       ,	UpperCamelCase:      Optional[int]=True       ,	UpperCamelCase:      str=32       ,	UpperCamelCase:      Any=5       ,	UpperCamelCase:      int=4       ,	UpperCamelCase:      List[Any]=37       ,	UpperCamelCase:      Optional[int]="gelu"       ,	UpperCamelCase:      Optional[Any]=0.1       ,	UpperCamelCase:      str=0.1       ,	UpperCamelCase:      List[Any]=10       ,	UpperCamelCase:      Any=0.02       ,	UpperCamelCase:      Any=[1, 16, 4, 4]       ,	UpperCamelCase:      Any=None       ,	)		->		Any:
											snake_case__									=     parent
											snake_case__									=     batch_size
											snake_case__									=     image_size
											snake_case__									=     patch_size
											snake_case__									=     num_channels
											snake_case__									=     is_training
											snake_case__									=     use_labels
											snake_case__									=     hidden_size
											snake_case__									=     num_hidden_layers
											snake_case__									=     num_attention_heads
											snake_case__									=     intermediate_size
											snake_case__									=     hidden_act
											snake_case__									=     hidden_dropout_prob
											snake_case__									=     attention_probs_dropout_prob
											snake_case__									=     type_sequence_label_size
											snake_case__									=     initializer_range
											snake_case__									=     scope
											snake_case__									=     backbone_featmap_shape

											# in ViT hybrid, the seq length equals the number of patches + 1 (we add 1 for the [CLS] token)
											# the number of patches is based on the feature map of the backbone, which by default uses an output stride
											# of 32, which means that the feature map has a spatial resolution of 1/32 of the input image size
											snake_case__									=     (self.image_size // 32) ** 2
											snake_case__									=     num_patches + 1
							def lowerCAmelCase_ (     self:      Dict						)		->		List[Any]:
											snake_case__									=     floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]						)

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

											snake_case__									=     self.get_config()

											return config, pixel_values, labels
							def lowerCAmelCase_ (     self:      Any						)		->		Optional[Any]:
											snake_case__									=     {
											    'global_padding': 'same',
											    'layer_type': 'bottleneck',
											    'depths': [3, 4, 9],
											    'out_features': ['stage1', 'stage2', 'stage3'],
											    'embedding_dynamic_padding': True,
											    'hidden_sizes': [4, 8, 16, 32],
											    'num_groups': 2,
											}

											return ViTHybridConfig(
											    image_size=self.image_size       ,	patch_size=self.patch_size       ,	num_channels=self.num_channels       ,	hidden_size=self.hidden_size       ,	num_hidden_layers=self.num_hidden_layers       ,	num_attention_heads=self.num_attention_heads       ,	intermediate_size=self.intermediate_size       ,	hidden_act=self.hidden_act       ,	hidden_dropout_prob=self.hidden_dropout_prob       ,	attention_probs_dropout_prob=self.attention_probs_dropout_prob       ,	is_decoder=UpperCamelCase       ,	initializer_range=self.initializer_range       ,	backbone_featmap_shape=self.backbone_featmap_shape       ,	backbone_config=UpperCamelCase       ,	)
							def lowerCAmelCase_ (     self:      List[str]       ,	UpperCamelCase:      str       ,	UpperCamelCase:      List[Any]       ,	UpperCamelCase:      List[Any]						)		->		Dict:
											snake_case__									=     ViTHybridModel(config=UpperCamelCase						)
											model.to(UpperCamelCase						)
											model.eval()
											snake_case__									=     model(UpperCamelCase						)
											self.parent.assertEqual(result.last_hidden_state.shape       ,	(self.batch_size, self.seq_length, self.hidden_size)						)
							def lowerCAmelCase_ (     self:      Dict       ,	UpperCamelCase:      Any       ,	UpperCamelCase:      str       ,	UpperCamelCase:      Dict						)		->		Tuple:
											snake_case__									=     self.type_sequence_label_size
											snake_case__									=     ViTHybridForImageClassification(UpperCamelCase						)
											model.to(UpperCamelCase						)
											model.eval()
											snake_case__									=     model(UpperCamelCase       ,	labels=UpperCamelCase						)
											self.parent.assertEqual(result.logits.shape       ,	(self.batch_size, self.type_sequence_label_size)						)


							def lowerCAmelCase_ (     self:      Any						)		->		List[str]:
											snake_case__									=     self.prepare_config_and_inputs()
											snake_case__							,						snake_case__							,						snake_case__									=     config_and_inputs
											snake_case__									=     {'pixel_values': pixel_values}
											return config, inputs_dict


@require_torch
class   __SCREAMING_SNAKE_CASE(	a_    ,  a_    ,  unittest.TestCase	):
							_UpperCAmelCase     =       (ViTHybridModel, ViTHybridForImageClassification) if is_torch_available() else ()
							_UpperCAmelCase     =       (
							    {"feature-extraction": ViTHybridModel, "image-classification": ViTHybridForImageClassification}
							    if is_torch_available()
							    else {}
							)
							_UpperCAmelCase     =       False
							_UpperCAmelCase     =       False
							_UpperCAmelCase     =       False
							def lowerCAmelCase_ (     self:      Optional[Any]						)		->		Tuple:
											snake_case__									=     ViTHybridModelTester(self						)
											snake_case__									=     ConfigTester(self       ,	config_class=UpperCamelCase       ,	has_text_modality=UpperCamelCase       ,	hidden_size=37						)
							def lowerCAmelCase_ (     self:      int						)		->		Union[str, Any]:
											self.config_tester.run_common_tests()
							@unittest.skip(reason='ViT does not use inputs_embeds'						)
							def lowerCAmelCase_ (     self:      Optional[int]						)		->		int:
											pass
							def lowerCAmelCase_ (     self:      Union[str, Any]						)		->		Dict:
											snake_case__							,						snake_case__									=     self.model_tester.prepare_config_and_inputs_for_common()

											for model_class in self.all_model_classes:
															snake_case__									=     model_class(UpperCamelCase						)
															self.assertIsInstance(model.get_input_embeddings()       ,	(nn.Module)						)
															snake_case__									=     model.get_output_embeddings()
															self.assertTrue(x is None or isinstance(UpperCamelCase       ,	nn.Linear						)						)
							def lowerCAmelCase_ (     self:      List[Any]						)		->		Any:
											snake_case__							,						snake_case__									=     self.model_tester.prepare_config_and_inputs_for_common()

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

															snake_case__									=     ['pixel_values']
															self.assertListEqual(arg_names[:1]       ,	UpperCamelCase						)
							def lowerCAmelCase_ (     self:      Dict						)		->		Tuple:
											snake_case__									=     self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_model(*UpperCamelCase						)
							def lowerCAmelCase_ (     self:      List[str]						)		->		int:
											snake_case__									=     self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_for_image_classification(*UpperCamelCase						)
							def lowerCAmelCase_ (     self:      Dict						)		->		Optional[int]:
											snake_case__							,						snake_case__									=     self.model_tester.prepare_config_and_inputs_for_common()

											snake_case__									=     _config_zero_init(UpperCamelCase						)
											for model_class in self.all_model_classes:
															snake_case__									=     model_class(config=UpperCamelCase						)
															# Skip the check for the backbone
															for name, module in model.named_modules():
																			if module.__class__.__name__ == "ViTHybridPatchEmbeddings":
																							snake_case__									=     [F'''{name}.{key}''' for key in module.state_dict().keys()]
																							break

															for name, param in model.named_parameters():
																			if param.requires_grad:
																							if name in backbone_params:
																											continue
																							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'''       ,	)


							@slow
							def lowerCAmelCase_ (     self:      Any						)		->		Dict:
											for model_name in VIT_HYBRID_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
															snake_case__									=     ViTHybridModel.from_pretrained(UpperCamelCase						)
															self.assertIsNotNone(UpperCamelCase						)


def      a_ (      )			->			Any:







				"""simple docstring"""
				snake_case__									=     Image.open('./tests/fixtures/tests_samples/COCO/000000039769.png'  )
				return image


@require_torch
@require_vision
class   __SCREAMING_SNAKE_CASE(	unittest.TestCase	):
							@cached_property
							def lowerCAmelCase_ (     self:      List[str]						)		->		Dict:
											return (
											    ViTHybridImageProcessor.from_pretrained(VIT_HYBRID_PRETRAINED_MODEL_ARCHIVE_LIST[0]						)
											    if is_vision_available()
											    else None
											)
							@slow
							def lowerCAmelCase_ (     self:      str						)		->		str:
											snake_case__									=     ViTHybridForImageClassification.from_pretrained(VIT_HYBRID_PRETRAINED_MODEL_ARCHIVE_LIST[0]						).to(
											    UpperCamelCase						)

											snake_case__									=     self.default_image_processor
											snake_case__									=     prepare_img()
											snake_case__									=     image_processor(images=UpperCamelCase       ,	return_tensors='pt'						).to(UpperCamelCase						)

											# forward pass
											with torch.no_grad():
															snake_case__									=     model(**UpperCamelCase						)

											# verify the logits
											snake_case__									=     torch.Size((1, 10_00)						)
											self.assertEqual(outputs.logits.shape       ,	UpperCamelCase						)

											snake_case__									=     torch.tensor([-1.9_090, -0.4_993, -0.2_389]						).to(UpperCamelCase						)

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


							@slow
							@require_accelerate
							def lowerCAmelCase_ (     self:      Any						)		->		Tuple:
											snake_case__									=     ViTHybridImageProcessor.from_pretrained('google/vit-hybrid-base-bit-384'						)
											snake_case__									=     ViTHybridForImageClassification.from_pretrained('google/vit-hybrid-base-bit-384'       ,	device_map='auto'						)

											snake_case__									=     prepare_img()

											snake_case__									=     image_processor(images=UpperCamelCase       ,	return_tensors='pt'						)
											snake_case__									=     model(**UpperCamelCase						)
											snake_case__									=     outputs.logits
											# model predicts one of the 1000 ImageNet classes
											snake_case__									=     logits.argmax(-1						).item()

											self.assertTrue(model.config.idalabel[predicted_class_idx]       ,	'tabby, tabby cat'						)











import gc
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    StableDiffusionSAGPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()


class   __SCREAMING_SNAKE_CASE(	a_    ,  a_    ,  unittest.TestCase	):
							_UpperCAmelCase     =       StableDiffusionSAGPipeline
							_UpperCAmelCase     =       TEXT_TO_IMAGE_PARAMS
							_UpperCAmelCase     =       TEXT_TO_IMAGE_BATCH_PARAMS
							_UpperCAmelCase     =       TEXT_TO_IMAGE_IMAGE_PARAMS
							_UpperCAmelCase     =       TEXT_TO_IMAGE_IMAGE_PARAMS
							_UpperCAmelCase     =       False
							def lowerCAmelCase_ (     self:      int						)		->		Any:
											torch.manual_seed(0						)
											snake_case__									=     UNetaDConditionModel(
											    block_out_channels=(32, 64)       ,	layers_per_block=2       ,	sample_size=32       ,	in_channels=4       ,	out_channels=4       ,	down_block_types=('DownBlock2D', 'CrossAttnDownBlock2D')       ,	up_block_types=('CrossAttnUpBlock2D', 'UpBlock2D')       ,	cross_attention_dim=32       ,	)
											snake_case__									=     DDIMScheduler(
											    beta_start=0.00_085       ,	beta_end=0.012       ,	beta_schedule='scaled_linear'       ,	clip_sample=UpperCamelCase       ,	set_alpha_to_one=UpperCamelCase       ,	)
											torch.manual_seed(0						)
											snake_case__									=     AutoencoderKL(
											    block_out_channels=[32, 64]       ,	in_channels=3       ,	out_channels=3       ,	down_block_types=['DownEncoderBlock2D', 'DownEncoderBlock2D']       ,	up_block_types=['UpDecoderBlock2D', 'UpDecoderBlock2D']       ,	latent_channels=4       ,	)
											torch.manual_seed(0						)
											snake_case__									=     CLIPTextConfig(
											    bos_token_id=0       ,	eos_token_id=2       ,	hidden_size=32       ,	intermediate_size=37       ,	layer_norm_eps=1e-05       ,	num_attention_heads=4       ,	num_hidden_layers=5       ,	pad_token_id=1       ,	vocab_size=10_00       ,	)
											snake_case__									=     CLIPTextModel(UpperCamelCase						)
											snake_case__									=     CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'						)

											snake_case__									=     {
											    'unet': unet,
											    'scheduler': scheduler,
											    'vae': vae,
											    'text_encoder': text_encoder,
											    'tokenizer': tokenizer,
											    'safety_checker': None,
											    'feature_extractor': None,
											}
											return components
							def lowerCAmelCase_ (     self:      Optional[int]       ,	UpperCamelCase:      Dict       ,	UpperCamelCase:      Optional[int]=0						)		->		Union[str, Any]:
											if str(UpperCamelCase						).startswith('mps'						):
															snake_case__									=     torch.manual_seed(UpperCamelCase						)
											else:
															snake_case__									=     torch.Generator(device=UpperCamelCase						).manual_seed(UpperCamelCase						)
											snake_case__									=     {
											    'prompt': '.',
											    'generator': generator,
											    'num_inference_steps': 2,
											    'guidance_scale': 1.0,
											    'sag_scale': 1.0,
											    'output_type': 'numpy',
											}
											return inputs


							def lowerCAmelCase_ (     self:      Union[str, Any]						)		->		Dict:
											super().test_inference_batch_single_identical(expected_max_diff=3e-3						)


@slow
@require_torch_gpu
class   __SCREAMING_SNAKE_CASE(	unittest.TestCase	):
							def lowerCAmelCase_ (     self:      Optional[int]						)		->		Optional[int]:
											# clean up the VRAM after each test
											super().tearDown()
											gc.collect()
											torch.cuda.empty_cache()
							def lowerCAmelCase_ (     self:      Dict						)		->		Tuple:
											snake_case__									=     StableDiffusionSAGPipeline.from_pretrained('CompVis/stable-diffusion-v1-4'						)
											snake_case__									=     sag_pipe.to(UpperCamelCase						)
											sag_pipe.set_progress_bar_config(disable=UpperCamelCase						)

											snake_case__									=     '.'
											snake_case__									=     torch.manual_seed(0						)
											snake_case__									=     sag_pipe(
											    [prompt]       ,	generator=UpperCamelCase       ,	guidance_scale=7.5       ,	sag_scale=1.0       ,	num_inference_steps=20       ,	output_type='np'						)

											snake_case__									=     output.images

											snake_case__									=     image[0, -3:, -3:, -1]

											assert image.shape == (1, 5_12, 5_12, 3)
											snake_case__									=     np.array([0.1_568, 0.1_738, 0.1_695, 0.1_693, 0.1_507, 0.1_705, 0.1_547, 0.1_751, 0.1_949]						)

											assert np.abs(image_slice.flatten() - expected_slice						).max() < 5e-2
							def lowerCAmelCase_ (     self:      Union[str, Any]						)		->		List[Any]:
											snake_case__									=     StableDiffusionSAGPipeline.from_pretrained('stabilityai/stable-diffusion-2-1-base'						)
											snake_case__									=     sag_pipe.to(UpperCamelCase						)
											sag_pipe.set_progress_bar_config(disable=UpperCamelCase						)

											snake_case__									=     '.'
											snake_case__									=     torch.manual_seed(0						)
											snake_case__									=     sag_pipe(
											    [prompt]       ,	generator=UpperCamelCase       ,	guidance_scale=7.5       ,	sag_scale=1.0       ,	num_inference_steps=20       ,	output_type='np'						)

											snake_case__									=     output.images

											snake_case__									=     image[0, -3:, -3:, -1]

											assert image.shape == (1, 5_12, 5_12, 3)
											snake_case__									=     np.array([0.3_459, 0.2_876, 0.2_537, 0.3_002, 0.2_671, 0.2_160, 0.3_026, 0.2_262, 0.2_371]						)

											assert np.abs(image_slice.flatten() - expected_slice						).max() < 5e-2


							def lowerCAmelCase_ (     self:      Union[str, Any]						)		->		Optional[Any]:
											snake_case__									=     StableDiffusionSAGPipeline.from_pretrained('stabilityai/stable-diffusion-2-1-base'						)
											snake_case__									=     sag_pipe.to(UpperCamelCase						)
											sag_pipe.set_progress_bar_config(disable=UpperCamelCase						)

											snake_case__									=     '.'
											snake_case__									=     torch.manual_seed(0						)
											snake_case__									=     sag_pipe(
											    [prompt]       ,	width=7_68       ,	height=5_12       ,	generator=UpperCamelCase       ,	guidance_scale=7.5       ,	sag_scale=1.0       ,	num_inference_steps=20       ,	output_type='np'       ,	)

											snake_case__									=     output.images

											assert image.shape == (1, 5_12, 7_68, 3)




