


import collections
import inspect
import unittest

from transformers import FocalNetConfig
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, _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 (
			    FocalNetBackbone,
			    FocalNetForImageClassification,
			    FocalNetForMaskedImageModeling,
			    FocalNetModel,
			)
			from transformers.models.focalnet.modeling_focalnet import FOCALNET_PRETRAINED_MODEL_ARCHIVE_LIST

if is_vision_available():
			from PIL import Image

			from transformers import AutoImageProcessor







class 				snake_case__:







						'''simple docstring'''




						def __init__(				self   ,						__lowercase   ,						__lowercase=1_3   ,						__lowercase=3_2   ,						__lowercase=2   ,						__lowercase=3   ,						__lowercase=1_6   ,						__lowercase=[3_2, 6_4, 1_2_8]   ,						__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=True   ,						__lowercase=1_0   ,						__lowercase=8   ,						__lowercase=["stage1", "stage2"]   ,						__lowercase=[1, 2]   ,						)					->    Union[str, Any]:
										lowerCAmelCase_ :      Union[str, Any]    				=				parent
										lowerCAmelCase_ :      Optional[int]    				=				batch_size
										lowerCAmelCase_ :      Union[str, Any]    				=				image_size
										lowerCAmelCase_ :      Tuple    				=				patch_size
										lowerCAmelCase_ :      Union[str, Any]    				=				num_channels
										lowerCAmelCase_ :      Optional[int]    				=				embed_dim
										lowerCAmelCase_ :      int    				=				hidden_sizes
										lowerCAmelCase_ :      Any    				=				depths
										lowerCAmelCase_ :      Optional[Any]    				=				num_heads
										lowerCAmelCase_ :      Optional[int]    				=				window_size
										lowerCAmelCase_ :      Union[str, Any]    				=				mlp_ratio
										lowerCAmelCase_ :      List[str]    				=				qkv_bias
										lowerCAmelCase_ :      Optional[Any]    				=				hidden_dropout_prob
										lowerCAmelCase_ :      Any    				=				attention_probs_dropout_prob
										lowerCAmelCase_ :      Optional[int]    				=				drop_path_rate
										lowerCAmelCase_ :      Optional[Any]    				=				hidden_act
										lowerCAmelCase_ :      str    				=				use_absolute_embeddings
										lowerCAmelCase_ :      Tuple    				=				patch_norm
										lowerCAmelCase_ :      Optional[int]    				=				layer_norm_eps
										lowerCAmelCase_ :      Tuple    				=				initializer_range
										lowerCAmelCase_ :      List[str]    				=				is_training
										lowerCAmelCase_ :      Tuple    				=				scope
										lowerCAmelCase_ :      Tuple    				=				use_labels
										lowerCAmelCase_ :      int    				=				type_sequence_label_size
										lowerCAmelCase_ :      Any    				=				encoder_stride
										lowerCAmelCase_ :      Dict    				=				out_features
										lowerCAmelCase_ :      List[str]    				=				out_indices




						def 						lowercase_ (				self							)					->    List[Any]:
										lowerCAmelCase_ :      str    				=				floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]							)

										lowerCAmelCase_ :      Dict    				=				None
										if self.use_labels:
														lowerCAmelCase_ :      Any    				=				ids_tensor([self.batch_size]   ,						self.type_sequence_label_size							)

										lowerCAmelCase_ :      Tuple    				=				self.get_config()

										return config, pixel_values, labels




						def 						lowercase_ (				self							)					->    Union[str, Any]:
										return FocalNetConfig(
										    image_size=self.image_size   ,						patch_size=self.patch_size   ,						num_channels=self.num_channels   ,						embed_dim=self.embed_dim   ,						hidden_sizes=self.hidden_sizes   ,						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   ,						encoder_stride=self.encoder_stride   ,						out_features=self.out_features   ,						out_indices=self.out_indices   ,						)




						def 						lowercase_ (				self   ,						__lowercase   ,						__lowercase   ,						__lowercase							)					->    List[str]:
										lowerCAmelCase_ :      Optional[Any]    				=				FocalNetModel(config=__lowercase							)
										model.to(__lowercase							)
										model.eval()
										lowerCAmelCase_ :      Dict    				=				model(__lowercase							)

										lowerCAmelCase_ :      str    				=				((config.image_size // config.patch_size) ** 2) // (4 ** (len(config.depths							) - 1))
										lowerCAmelCase_ :      int    				=				int(config.embed_dim * 2 ** (len(config.depths							) - 1)							)

										self.parent.assertEqual(result.last_hidden_state.shape   ,						(self.batch_size, expected_seq_len, expected_dim)							)




						def 						lowercase_ (				self   ,						__lowercase   ,						__lowercase   ,						__lowercase							)					->    Tuple:
										lowerCAmelCase_ :      str    				=				FocalNetBackbone(config=__lowercase							)
										model.to(__lowercase							)
										model.eval()
										lowerCAmelCase_ :      str    				=				model(__lowercase							)

										# 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.image_size, 8, 8]							)

										# 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
										lowerCAmelCase_ :      Optional[int]    				=				None
										lowerCAmelCase_ :      Dict    				=				FocalNetBackbone(config=__lowercase							)
										model.to(__lowercase							)
										model.eval()
										lowerCAmelCase_ :      str    				=				model(__lowercase							)

										# verify feature maps
										self.parent.assertEqual(len(result.feature_maps							)   ,						1							)
										self.parent.assertListEqual(list(result.feature_maps[0].shape							)   ,						[self.batch_size, self.image_size * 2, 4, 4]							)

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




						def 						lowercase_ (				self   ,						__lowercase   ,						__lowercase   ,						__lowercase							)					->    Tuple:
										lowerCAmelCase_ :      Dict    				=				FocalNetForMaskedImageModeling(config=__lowercase							)
										model.to(__lowercase							)
										model.eval()
										lowerCAmelCase_ :      str    				=				model(__lowercase							)
										self.parent.assertEqual(
										    result.reconstruction.shape   ,						(self.batch_size, self.num_channels, self.image_size, self.image_size)							)

										# test greyscale images
										lowerCAmelCase_ :      Union[str, Any]    				=				1
										lowerCAmelCase_ :      Optional[int]    				=				FocalNetForMaskedImageModeling(__lowercase							)
										model.to(__lowercase							)
										model.eval()

										lowerCAmelCase_ :      Optional[int]    				=				floats_tensor([self.batch_size, 1, self.image_size, self.image_size]							)
										lowerCAmelCase_ :      Any    				=				model(__lowercase							)
										self.parent.assertEqual(result.reconstruction.shape   ,						(self.batch_size, 1, self.image_size, self.image_size)							)




						def 						lowercase_ (				self   ,						__lowercase   ,						__lowercase   ,						__lowercase							)					->    Dict:
										lowerCAmelCase_ :      Optional[Any]    				=				self.type_sequence_label_size
										lowerCAmelCase_ :      Dict    				=				FocalNetForImageClassification(__lowercase							)
										model.to(__lowercase							)
										model.eval()
										lowerCAmelCase_ :      Optional[Any]    				=				model(__lowercase   ,						labels=__lowercase							)
										self.parent.assertEqual(result.logits.shape   ,						(self.batch_size, self.type_sequence_label_size)							)

										# test greyscale images
										lowerCAmelCase_ :      Optional[int]    				=				1
										lowerCAmelCase_ :      List[str]    				=				FocalNetForImageClassification(__lowercase							)
										model.to(__lowercase							)
										model.eval()

										lowerCAmelCase_ :      List[Any]    				=				floats_tensor([self.batch_size, 1, self.image_size, self.image_size]							)
										lowerCAmelCase_ :      Any    				=				model(__lowercase							)
										self.parent.assertEqual(result.logits.shape   ,						(self.batch_size, self.type_sequence_label_size)							)




						def 						lowercase_ (				self							)					->    Dict:
										lowerCAmelCase_ :      Optional[Any]    				=				self.prepare_config_and_inputs()

										lowerCAmelCase_ :      List[str]    				=				config_and_inputs
										lowerCAmelCase_ :      Optional[Any]    				=				{'''pixel_values''': pixel_values}
										return config, inputs_dict







@require_torch
class 				snake_case__(							UpperCAmelCase__,						UpperCAmelCase__,						unittest.TestCase							):







						'''simple docstring'''




						SCREAMING_SNAKE_CASE__      :  Any					=							(
						    (
						        FocalNetModel,
						        FocalNetForImageClassification,
						        FocalNetForMaskedImageModeling,
						        FocalNetBackbone,
						    )
						    if is_torch_available()
						    else ()
						)
						SCREAMING_SNAKE_CASE__      :  Union[str, Any]					=							(
						    {"""feature-extraction""": FocalNetModel, """image-classification""": FocalNetForImageClassification}
						    if is_torch_available()
						    else {}
						)
						SCREAMING_SNAKE_CASE__      :  Tuple					=							False

						SCREAMING_SNAKE_CASE__      :  Tuple					=							False
						SCREAMING_SNAKE_CASE__      :  Union[str, Any]					=							False
						SCREAMING_SNAKE_CASE__      :  Union[str, Any]					=							False
						SCREAMING_SNAKE_CASE__      :  Optional[int]					=							False




						def 						lowercase_ (				self							)					->    Optional[Any]:
										lowerCAmelCase_ :      Any    				=				FocalNetModelTester(self							)
										lowerCAmelCase_ :      Union[str, Any]    				=				ConfigTester(self   ,						config_class=__lowercase   ,						embed_dim=3_7   ,						has_text_modality=__lowercase							)




						def 						lowercase_ (				self							)					->    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 						lowercase_ (				self							)					->    Any:
										return




						def 						lowercase_ (				self							)					->    List[str]:
										lowerCAmelCase_ :      str    				=				self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_model(*__lowercase							)




						def 						lowercase_ (				self							)					->    Union[str, Any]:
										lowerCAmelCase_ :      Tuple    				=				self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_backbone(*__lowercase							)




						def 						lowercase_ (				self							)					->    List[Any]:
										lowerCAmelCase_ :      List[str]    				=				self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_for_masked_image_modeling(*__lowercase							)




						def 						lowercase_ (				self							)					->    str:
										lowerCAmelCase_ :      List[str]    				=				self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_for_image_classification(*__lowercase							)




						@unittest.skip(reason='''FocalNet does not use inputs_embeds'''							)
						def 						lowercase_ (				self							)					->    List[Any]:
										pass




						@unittest.skip(reason='''FocalNet does not use feedforward chunking'''							)
						def 						lowercase_ (				self							)					->    Any:
										pass




						def 						lowercase_ (				self							)					->    Any:
										lowerCAmelCase_ :      Optional[Any]    				=				self.model_tester.prepare_config_and_inputs_for_common()

										for model_class in self.all_model_classes[:-1]:
														lowerCAmelCase_ :      int    				=				model_class(__lowercase							)
														self.assertIsInstance(model.get_input_embeddings()   ,						(nn.Module)							)
														lowerCAmelCase_ :      Union[str, Any]    				=				model.get_output_embeddings()
														self.assertTrue(x is None or isinstance(__lowercase   ,						nn.Linear							)							)




						def 						lowercase_ (				self							)					->    List[Any]:
										lowerCAmelCase_ :      int    				=				self.model_tester.prepare_config_and_inputs_for_common()

										for model_class in self.all_model_classes[:-1]:
														lowerCAmelCase_ :      Optional[int]    				=				model_class(__lowercase							)
														lowerCAmelCase_ :      Tuple    				=				inspect.signature(model.forward							)
														# signature.parameters is an OrderedDict => so arg_names order is deterministic
														lowerCAmelCase_ :      Any    				=				[*signature.parameters.keys()]

														lowerCAmelCase_ :      Optional[int]    				=				['''pixel_values''']
														self.assertListEqual(arg_names[:1]   ,						__lowercase							)




						def 						lowercase_ (				self   ,						__lowercase   ,						__lowercase   ,						__lowercase   ,						__lowercase							)					->    Tuple:
										lowerCAmelCase_ :      str    				=				model_class(__lowercase							)
										model.to(__lowercase							)
										model.eval()

										with torch.no_grad():
														lowerCAmelCase_ :      Optional[int]    				=				model(**self._prepare_for_class(__lowercase   ,						__lowercase							)							)

										lowerCAmelCase_ :      Tuple    				=				outputs.hidden_states

										lowerCAmelCase_ :      Optional[Any]    				=				getattr(
										    self.model_tester   ,						'''expected_num_hidden_layers'''   ,						len(self.model_tester.depths							) + 1							)
										self.assertEqual(len(__lowercase							)   ,						__lowercase							)

										# FocalNet has a different seq_length
										lowerCAmelCase_ :      List[Any]    				=				(
										    config.patch_size
										    if isinstance(config.patch_size   ,						collections.abc.Iterable							)
										    else (config.patch_size, config.patch_size)
										)

										lowerCAmelCase_ :      Any    				=				(image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0])

										self.assertListEqual(
										    list(hidden_states[0].shape[-2:]							)   ,						[num_patches, self.model_tester.embed_dim]   ,						)

										lowerCAmelCase_ :      Any    				=				outputs.reshaped_hidden_states
										self.assertEqual(len(__lowercase							)   ,						__lowercase							)

										lowerCAmelCase_ :      List[Any]    				=				reshaped_hidden_states[0].shape
										lowerCAmelCase_ :      Tuple    				=				(
										    reshaped_hidden_states[0].view(__lowercase   ,						__lowercase   ,						height * width							).permute(0   ,						2   ,						1							)
										)
										self.assertListEqual(
										    list(reshaped_hidden_states.shape[-2:]							)   ,						[num_patches, self.model_tester.embed_dim]   ,						)




						def 						lowercase_ (				self							)					->    int:
										lowerCAmelCase_ :      List[str]    				=				self.model_tester.prepare_config_and_inputs_for_common()

										lowerCAmelCase_ :      Union[str, Any]    				=				(
										    self.model_tester.image_size
										    if isinstance(self.model_tester.image_size   ,						collections.abc.Iterable							)
										    else (self.model_tester.image_size, self.model_tester.image_size)
										)

										for model_class in self.all_model_classes[:-1]:
														lowerCAmelCase_ :      Optional[int]    				=				True
														self.check_hidden_states_output(__lowercase   ,						__lowercase   ,						__lowercase   ,						__lowercase							)

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

														self.check_hidden_states_output(__lowercase   ,						__lowercase   ,						__lowercase   ,						__lowercase							)




						def 						lowercase_ (				self							)					->    Any:
										lowerCAmelCase_ :      str    				=				self.model_tester.prepare_config_and_inputs_for_common()
										lowerCAmelCase_ :      Optional[int]    				=				3

										lowerCAmelCase_ :      List[str]    				=				(
										    self.model_tester.image_size
										    if isinstance(self.model_tester.image_size   ,						collections.abc.Iterable							)
										    else (self.model_tester.image_size, self.model_tester.image_size)
										)
										lowerCAmelCase_ :      Optional[int]    				=				(
										    config.patch_size
										    if isinstance(config.patch_size   ,						collections.abc.Iterable							)
										    else (config.patch_size, config.patch_size)
										)

										lowerCAmelCase_ :      Optional[Any]    				=				image_size[0] + patch_size[0] - (image_size[0] % patch_size[0])
										lowerCAmelCase_ :      List[Any]    				=				image_size[1] + patch_size[1] - (image_size[1] % patch_size[1])

										for model_class in self.all_model_classes[:-1]:
														lowerCAmelCase_ :      List[Any]    				=				True
														self.check_hidden_states_output(__lowercase   ,						__lowercase   ,						__lowercase   ,						(padded_height, padded_width)							)

														# check that output_hidden_states also work using config
														del inputs_dict["output_hidden_states"]
														lowerCAmelCase_ :      Optional[Any]    				=				True
														self.check_hidden_states_output(__lowercase   ,						__lowercase   ,						__lowercase   ,						(padded_height, padded_width)							)




						@slow
						def 						lowercase_ (				self							)					->    Dict:
										for model_name in FOCALNET_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
														lowerCAmelCase_ :      Tuple    				=				FocalNetModel.from_pretrained(__lowercase							)
														self.assertIsNotNone(__lowercase							)




						def 						lowercase_ (				self							)					->    str:
										lowerCAmelCase_ :      Any    				=				self.model_tester.prepare_config_and_inputs_for_common()

										lowerCAmelCase_ :      Optional[int]    				=				_config_zero_init(__lowercase							)
										for model_class in self.all_model_classes:
														lowerCAmelCase_ :      Dict    				=				model_class(config=__lowercase							)
														for name, param in model.named_parameters():
																		if "embeddings" not in name and 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"""   ,						)







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







						'''simple docstring'''




						@cached_property
						def 						lowercase_ (				self							)					->    int:
										# TODO update organization
										return AutoImageProcessor.from_pretrained('''microsoft/focalnet-tiny'''							) if is_vision_available() else None




						@slow
						def 						lowercase_ (				self							)					->    Optional[int]:
										lowerCAmelCase_ :      Optional[int]    				=				FocalNetForImageClassification.from_pretrained('''microsoft/focalnet-tiny'''							).to(__lowercase							)
										lowerCAmelCase_ :      Any    				=				self.default_image_processor

										lowerCAmelCase_ :      Dict    				=				Image.open('''./tests/fixtures/tests_samples/COCO/000000039769.png'''							)
										lowerCAmelCase_ :      Any    				=				image_processor(images=__lowercase   ,						return_tensors='''pt'''							).to(__lowercase							)

										# forward pass
										with torch.no_grad():
														lowerCAmelCase_ :      Tuple    				=				model(**__lowercase							)

										# verify the logits
										lowerCAmelCase_ :      Optional[Any]    				=				torch.Size((1, 1_0_0_0)							)
										self.assertEqual(outputs.logits.shape   ,						__lowercase							)
										lowerCAmelCase_ :      int    				=				torch.tensor([0.21_66, -0.43_68, 0.21_91]							).to(__lowercase							)
										self.assertTrue(torch.allclose(outputs.logits[0, :3]   ,						__lowercase   ,						atol=1e-4							)							)
										self.assertTrue(outputs.logits.argmax(dim=-1							).item()   ,						2_8_1							)







@require_torch
class 				snake_case__(							UpperCAmelCase__,						unittest.TestCase							):







						'''simple docstring'''




						SCREAMING_SNAKE_CASE__      :  List[Any]					=							(FocalNetBackbone,) if is_torch_available() else ()
						SCREAMING_SNAKE_CASE__      :  Optional[Any]					=							FocalNetConfig

						SCREAMING_SNAKE_CASE__      :  str					=							False




						def 						lowercase_ (				self							)					->    Union[str, Any]:
										lowerCAmelCase_ :      Tuple    				=				FocalNetModelTester(self							)



import gc
import unittest

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

from diffusers import TransformeraDModel, VQDiffusionPipeline, VQDiffusionScheduler, VQModel
from diffusers.pipelines.vq_diffusion.pipeline_vq_diffusion import LearnedClassifierFreeSamplingEmbeddings
from diffusers.utils import load_numpy, slow, torch_device
from diffusers.utils.testing_utils import require_torch_gpu


_UpperCAmelCase			:  Any     =False







class 				snake_case__(							unittest.TestCase							):







						'''simple docstring'''




						def 						lowercase_ (				self							)					->    Optional[int]:
										# clean up the VRAM after each test
										super().tearDown()
										gc.collect()
										torch.cuda.empty_cache()




						@property
						def 						lowercase_ (				self							)					->    Union[str, Any]:
										return 1_2




						@property
						def 						lowercase_ (				self							)					->    Any:
										return 1_2




						@property
						def 						lowercase_ (				self							)					->    Optional[Any]:
										return 3_2




						@property
						def 						lowercase_ (				self							)					->    int:
										torch.manual_seed(0							)
										lowerCAmelCase_ :      Any    				=				VQModel(
										    block_out_channels=[3_2, 6_4]   ,						in_channels=3   ,						out_channels=3   ,						down_block_types=['''DownEncoderBlock2D''', '''DownEncoderBlock2D''']   ,						up_block_types=['''UpDecoderBlock2D''', '''UpDecoderBlock2D''']   ,						latent_channels=3   ,						num_vq_embeddings=self.num_embed   ,						vq_embed_dim=3   ,						)
										return model




						@property
						def 						lowercase_ (				self							)					->    Dict:
										lowerCAmelCase_ :      Optional[Any]    				=				CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''							)
										return tokenizer




						@property
						def 						lowercase_ (				self							)					->    int:
										torch.manual_seed(0							)
										lowerCAmelCase_ :      Union[str, Any]    				=				CLIPTextConfig(
										    bos_token_id=0   ,						eos_token_id=2   ,						hidden_size=self.text_embedder_hidden_size   ,						intermediate_size=3_7   ,						layer_norm_eps=1e-05   ,						num_attention_heads=4   ,						num_hidden_layers=5   ,						pad_token_id=1   ,						vocab_size=1_0_0_0   ,						)
										return CLIPTextModel(__lowercase							)




						@property
						def 						lowercase_ (				self							)					->    List[str]:
										torch.manual_seed(0							)

										lowerCAmelCase_ :      Union[str, Any]    				=				1_2
										lowerCAmelCase_ :      int    				=				1_2

										lowerCAmelCase_ :      Union[str, Any]    				=				{
										    '''attention_bias''': True,
										    '''cross_attention_dim''': 3_2,
										    '''attention_head_dim''': height * width,
										    '''num_attention_heads''': 1,
										    '''num_vector_embeds''': self.num_embed,
										    '''num_embeds_ada_norm''': self.num_embeds_ada_norm,
										    '''norm_num_groups''': 3_2,
										    '''sample_size''': width,
										    '''activation_fn''': '''geglu-approximate''',
										}

										lowerCAmelCase_ :      List[str]    				=				TransformeraDModel(**__lowercase							)
										return model




						def 						lowercase_ (				self							)					->    str:
										lowerCAmelCase_ :      List[Any]    				=				'''cpu'''

										lowerCAmelCase_ :      Any    				=				self.dummy_vqvae
										lowerCAmelCase_ :      str    				=				self.dummy_text_encoder
										lowerCAmelCase_ :      Union[str, Any]    				=				self.dummy_tokenizer
										lowerCAmelCase_ :      int    				=				self.dummy_transformer
										lowerCAmelCase_ :      List[str]    				=				VQDiffusionScheduler(self.num_embed							)
										lowerCAmelCase_ :      Union[str, Any]    				=				LearnedClassifierFreeSamplingEmbeddings(learnable=__lowercase							)

										lowerCAmelCase_ :      Dict    				=				VQDiffusionPipeline(
										    vqvae=__lowercase   ,						text_encoder=__lowercase   ,						tokenizer=__lowercase   ,						transformer=__lowercase   ,						scheduler=__lowercase   ,						learned_classifier_free_sampling_embeddings=__lowercase   ,						)
										lowerCAmelCase_ :      int    				=				pipe.to(__lowercase							)
										pipe.set_progress_bar_config(disable=__lowercase							)

										lowerCAmelCase_ :      Any    				=				'''teddy bear playing in the pool'''

										lowerCAmelCase_ :      int    				=				torch.Generator(device=__lowercase							).manual_seed(0							)
										lowerCAmelCase_ :      Tuple    				=				pipe([prompt]   ,						generator=__lowercase   ,						num_inference_steps=2   ,						output_type='''np'''							)
										lowerCAmelCase_ :      Union[str, Any]    				=				output.images

										lowerCAmelCase_ :      List[Any]    				=				torch.Generator(device=__lowercase							).manual_seed(0							)
										lowerCAmelCase_ :      List[Any]    				=				pipe(
										    [prompt]   ,						generator=__lowercase   ,						output_type='''np'''   ,						return_dict=__lowercase   ,						num_inference_steps=2							)[0]

										lowerCAmelCase_ :      List[str]    				=				image[0, -3:, -3:, -1]
										lowerCAmelCase_ :      str    				=				image_from_tuple[0, -3:, -3:, -1]

										assert image.shape == (1, 2_4, 2_4, 3)

										lowerCAmelCase_ :      Optional[int]    				=				np.array([0.65_51, 0.61_68, 0.50_08, 0.56_76, 0.56_59, 0.42_95, 0.60_73, 0.55_99, 0.49_92]							)

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




						def 						lowercase_ (				self							)					->    List[str]:
										lowerCAmelCase_ :      Optional[Any]    				=				'''cpu'''

										lowerCAmelCase_ :      str    				=				self.dummy_vqvae
										lowerCAmelCase_ :      Dict    				=				self.dummy_text_encoder
										lowerCAmelCase_ :      List[Any]    				=				self.dummy_tokenizer
										lowerCAmelCase_ :      Union[str, Any]    				=				self.dummy_transformer
										lowerCAmelCase_ :      Tuple    				=				VQDiffusionScheduler(self.num_embed							)
										lowerCAmelCase_ :      str    				=				LearnedClassifierFreeSamplingEmbeddings(
										    learnable=__lowercase   ,						hidden_size=self.text_embedder_hidden_size   ,						length=tokenizer.model_max_length							)

										lowerCAmelCase_ :      List[str]    				=				VQDiffusionPipeline(
										    vqvae=__lowercase   ,						text_encoder=__lowercase   ,						tokenizer=__lowercase   ,						transformer=__lowercase   ,						scheduler=__lowercase   ,						learned_classifier_free_sampling_embeddings=__lowercase   ,						)
										lowerCAmelCase_ :      Union[str, Any]    				=				pipe.to(__lowercase							)
										pipe.set_progress_bar_config(disable=__lowercase							)

										lowerCAmelCase_ :      Any    				=				'''teddy bear playing in the pool'''

										lowerCAmelCase_ :      List[str]    				=				torch.Generator(device=__lowercase							).manual_seed(0							)
										lowerCAmelCase_ :      Dict    				=				pipe([prompt]   ,						generator=__lowercase   ,						num_inference_steps=2   ,						output_type='''np'''							)
										lowerCAmelCase_ :      str    				=				output.images

										lowerCAmelCase_ :      List[Any]    				=				torch.Generator(device=__lowercase							).manual_seed(0							)
										lowerCAmelCase_ :      Union[str, Any]    				=				pipe(
										    [prompt]   ,						generator=__lowercase   ,						output_type='''np'''   ,						return_dict=__lowercase   ,						num_inference_steps=2							)[0]

										lowerCAmelCase_ :      List[str]    				=				image[0, -3:, -3:, -1]
										lowerCAmelCase_ :      str    				=				image_from_tuple[0, -3:, -3:, -1]

										assert image.shape == (1, 2_4, 2_4, 3)

										lowerCAmelCase_ :      Union[str, Any]    				=				np.array([0.66_93, 0.60_75, 0.49_59, 0.57_01, 0.55_83, 0.43_33, 0.61_71, 0.56_84, 0.49_88]							)

										assert np.abs(image_slice.flatten() - expected_slice							).max() < 2.0
										assert np.abs(image_from_tuple_slice.flatten() - expected_slice							).max() < 1e-2







@slow
@require_torch_gpu
class 				snake_case__(							unittest.TestCase							):







						'''simple docstring'''




						def 						lowercase_ (				self							)					->    Optional[int]:
										# clean up the VRAM after each test
										super().tearDown()
										gc.collect()
										torch.cuda.empty_cache()




						def 						lowercase_ (				self							)					->    int:
										lowerCAmelCase_ :      Tuple    				=				load_numpy(
										    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
										    '''/vq_diffusion/teddy_bear_pool_classifier_free_sampling.npy'''							)

										lowerCAmelCase_ :      str    				=				VQDiffusionPipeline.from_pretrained('''microsoft/vq-diffusion-ithq'''							)
										lowerCAmelCase_ :      List[Any]    				=				pipeline.to(__lowercase							)
										pipeline.set_progress_bar_config(disable=__lowercase							)

										# requires GPU generator for gumbel softmax
										# don't use GPU generator in tests though
										lowerCAmelCase_ :      List[Any]    				=				torch.Generator(device=__lowercase							).manual_seed(0							)
										lowerCAmelCase_ :      Optional[int]    				=				pipeline(
										    '''teddy bear playing in the pool'''   ,						num_images_per_prompt=1   ,						generator=__lowercase   ,						output_type='''np'''   ,						)

										lowerCAmelCase_ :      Union[str, Any]    				=				output.images[0]

										assert image.shape == (2_5_6, 2_5_6, 3)
										assert np.abs(expected_image - image							).max() < 2.0