



"""simple docstring"""





from typing import Optional, Tuple, Union

import flax
import flax.linen as nn
import jax
import jax.numpy as jnp
from flax.core.frozen_dict import FrozenDict

from ..configuration_utils import ConfigMixin, flax_register_to_config
from ..utils import BaseOutput
from .embeddings_flax import FlaxTimestepEmbedding, FlaxTimesteps
from .modeling_flax_utils import FlaxModelMixin
from .unet_ad_blocks_flax import (
    FlaxCrossAttnDownBlockaD,
    FlaxCrossAttnUpBlockaD,
    FlaxDownBlockaD,
    FlaxUNetMidBlockaDCrossAttn,
    FlaxUpBlockaD,
)




@flax.struct.dataclass
class 			lowercase__			(      A    ):
							'''simple docstring'''







							_UpperCAmelCase      				=					42







@flax_register_to_config
class 			lowercase__			(      nn.Module,					A,					A    ):
							'''simple docstring'''







							_UpperCAmelCase      				=					32
							_UpperCAmelCase      				=					4
							_UpperCAmelCase      				=					4
							_UpperCAmelCase      				=					(
							    "CrossAttnDownBlock2D",
							    "CrossAttnDownBlock2D",
							    "CrossAttnDownBlock2D",
							    "DownBlock2D",
							)
							_UpperCAmelCase      				=					("UpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D")
							_UpperCAmelCase      				=					False
							_UpperCAmelCase      				=					(3_20, 6_40, 12_80, 12_80)
							_UpperCAmelCase      				=					2
							_UpperCAmelCase      				=					8
							_UpperCAmelCase      				=					None
							_UpperCAmelCase      				=					12_80
							_UpperCAmelCase      				=					0.0
							_UpperCAmelCase      				=					False
							_UpperCAmelCase      				=					jnp.floataa
							_UpperCAmelCase      				=					True
							_UpperCAmelCase      				=					0
							_UpperCAmelCase      				=					False






							def        lowerCamelCase_   (  self	,       snake_case						)			->							FrozenDict:
									# init input tensors
									_UpperCAmelCase							 =		(1, self.in_channels, self.sample_size, self.sample_size)
									_UpperCAmelCase							 =		jnp.zeros(snake_case	,       dtype=jnp.floataa						)
									_UpperCAmelCase							 =		jnp.ones((1,)	,       dtype=jnp.intaa						)
									_UpperCAmelCase							 =		jnp.zeros((1, 1, self.cross_attention_dim)	,       dtype=jnp.floataa						)

									_UpperCAmelCase ,					_UpperCAmelCase							 =		jax.random.split(snake_case						)
									_UpperCAmelCase							 =		{'params': params_rng, 'dropout': dropout_rng}

									return self.init(snake_case	,       snake_case	,       snake_case	,       snake_case						)["params"]






							def        lowerCamelCase_   (  self						)			->							List[str]:
									_UpperCAmelCase							 =		self.block_out_channels
									_UpperCAmelCase							 =		block_out_channels[0] * 4

									if self.num_attention_heads is not None:
											raise ValueError(
											    'At the moment it is not possible to define the number of attention heads via `num_attention_heads` because of a naming issue as described in https://github.com/huggingface/diffusers/issues/2011#issuecomment-1547958131. Passing `num_attention_heads` will only be supported in diffusers v0.19.'						)

									# If `num_attention_heads` is not defined (which is the case for most models)
									# it will default to `attention_head_dim`. This looks weird upon first reading it and it is.
									# The reason for this behavior is to correct for incorrectly named variables that were introduced
									# when this library was created. The incorrect naming was only discovered much later in https://github.com/huggingface/diffusers/issues/2011#issuecomment-1547958131
									# Changing `attention_head_dim` to `num_attention_heads` for 40,000+ configurations is too backwards breaking
									# which is why we correct for the naming here.
									_UpperCAmelCase							 =		self.num_attention_heads or self.attention_head_dim

									# input
									_UpperCAmelCase							 =		nn.Conv(
									    block_out_channels[0]	,       kernel_size=(3, 3)	,       strides=(1, 1)	,       padding=((1, 1), (1, 1))	,       dtype=self.dtype	,       )

									# time
									_UpperCAmelCase							 =		FlaxTimesteps(
									    block_out_channels[0]	,       flip_sin_to_cos=self.flip_sin_to_cos	,       freq_shift=self.config.freq_shift						)
									_UpperCAmelCase							 =		FlaxTimestepEmbedding(snake_case	,       dtype=self.dtype						)

									_UpperCAmelCase							 =		self.only_cross_attention
									if isinstance(snake_case	,       snake_case						):
											_UpperCAmelCase							 =		(only_cross_attention,) * len(self.down_block_types						)

									if isinstance(snake_case	,       snake_case						):
											_UpperCAmelCase							 =		(num_attention_heads,) * len(self.down_block_types						)

									# down
									_UpperCAmelCase							 =		[]
									_UpperCAmelCase							 =		block_out_channels[0]
									for i, down_block_type in enumerate(self.down_block_types						):
											_UpperCAmelCase							 =		output_channel
											_UpperCAmelCase							 =		block_out_channels[i]
											_UpperCAmelCase							 =		i == len(snake_case						) - 1

											if down_block_type == "CrossAttnDownBlock2D":
													_UpperCAmelCase							 =		FlaxCrossAttnDownBlockaD(
													    in_channels=snake_case	,       out_channels=snake_case	,       dropout=self.dropout	,       num_layers=self.layers_per_block	,       num_attention_heads=num_attention_heads[i]	,       add_downsample=not is_final_block	,       use_linear_projection=self.use_linear_projection	,       only_cross_attention=only_cross_attention[i]	,       use_memory_efficient_attention=self.use_memory_efficient_attention	,       dtype=self.dtype	,       )
											else:
													_UpperCAmelCase							 =		FlaxDownBlockaD(
													    in_channels=snake_case	,       out_channels=snake_case	,       dropout=self.dropout	,       num_layers=self.layers_per_block	,       add_downsample=not is_final_block	,       dtype=self.dtype	,       )

											down_blocks.append(snake_case						)
									_UpperCAmelCase							 =		down_blocks

									# mid
									_UpperCAmelCase							 =		FlaxUNetMidBlockaDCrossAttn(
									    in_channels=block_out_channels[-1]	,       dropout=self.dropout	,       num_attention_heads=num_attention_heads[-1]	,       use_linear_projection=self.use_linear_projection	,       use_memory_efficient_attention=self.use_memory_efficient_attention	,       dtype=self.dtype	,       )

									# up
									_UpperCAmelCase							 =		[]
									_UpperCAmelCase							 =		list(reversed(snake_case						)						)
									_UpperCAmelCase							 =		list(reversed(snake_case						)						)
									_UpperCAmelCase							 =		list(reversed(snake_case						)						)
									_UpperCAmelCase							 =		reversed_block_out_channels[0]
									for i, up_block_type in enumerate(self.up_block_types						):
											_UpperCAmelCase							 =		output_channel
											_UpperCAmelCase							 =		reversed_block_out_channels[i]
											_UpperCAmelCase							 =		reversed_block_out_channels[min(i + 1	,       len(snake_case						) - 1						)]

											_UpperCAmelCase							 =		i == len(snake_case						) - 1

											if up_block_type == "CrossAttnUpBlock2D":
													_UpperCAmelCase							 =		FlaxCrossAttnUpBlockaD(
													    in_channels=snake_case	,       out_channels=snake_case	,       prev_output_channel=snake_case	,       num_layers=self.layers_per_block + 1	,       num_attention_heads=reversed_num_attention_heads[i]	,       add_upsample=not is_final_block	,       dropout=self.dropout	,       use_linear_projection=self.use_linear_projection	,       only_cross_attention=only_cross_attention[i]	,       use_memory_efficient_attention=self.use_memory_efficient_attention	,       dtype=self.dtype	,       )
											else:
													_UpperCAmelCase							 =		FlaxUpBlockaD(
													    in_channels=snake_case	,       out_channels=snake_case	,       prev_output_channel=snake_case	,       num_layers=self.layers_per_block + 1	,       add_upsample=not is_final_block	,       dropout=self.dropout	,       dtype=self.dtype	,       )

											up_blocks.append(snake_case						)
											_UpperCAmelCase							 =		output_channel
									_UpperCAmelCase							 =		up_blocks

									# out
									_UpperCAmelCase							 =		nn.GroupNorm(num_groups=32	,       epsilon=1E-5						)
									_UpperCAmelCase							 =		nn.Conv(
									    self.out_channels	,       kernel_size=(3, 3)	,       strides=(1, 1)	,       padding=((1, 1), (1, 1))	,       dtype=self.dtype	,       )






							def __call__(  self	,       snake_case	,       snake_case	,       snake_case	,       snake_case=None	,       snake_case=None	,       snake_case = True	,       snake_case = False	,       )			->							Union[FlaxUNetaDConditionOutput, Tuple]:
									# 1. time
									if not isinstance(snake_case	,       jnp.ndarray						):
											_UpperCAmelCase							 =		jnp.array([timesteps]	,       dtype=jnp.intaa						)
									elif isinstance(snake_case	,       jnp.ndarray						) and len(timesteps.shape						) == 0:
											_UpperCAmelCase							 =		timesteps.astype(dtype=jnp.floataa						)
											_UpperCAmelCase							 =		jnp.expand_dims(snake_case	,       0						)

									_UpperCAmelCase							 =		self.time_proj(snake_case						)
									_UpperCAmelCase							 =		self.time_embedding(snake_case						)

									# 2. pre-process
									_UpperCAmelCase							 =		jnp.transpose(snake_case	,       (0, 2, 3, 1)						)
									_UpperCAmelCase							 =		self.conv_in(snake_case						)

									# 3. down
									_UpperCAmelCase							 =		(sample,)
									for down_block in self.down_blocks:
											if isinstance(snake_case	,       snake_case						):
													_UpperCAmelCase ,					_UpperCAmelCase							 =		down_block(snake_case	,       snake_case	,       snake_case	,       deterministic=not train						)
											else:
													_UpperCAmelCase ,					_UpperCAmelCase							 =		down_block(snake_case	,       snake_case	,       deterministic=not train						)
											down_block_res_samples += res_samples

									if down_block_additional_residuals is not None:
											_UpperCAmelCase							 =		()

											for down_block_res_sample, down_block_additional_residual in zip(
											    snake_case	,       snake_case						):
													down_block_res_sample += down_block_additional_residual
													new_down_block_res_samples += (down_block_res_sample,)

											_UpperCAmelCase							 =		new_down_block_res_samples

									# 4. mid
									_UpperCAmelCase							 =		self.mid_block(snake_case	,       snake_case	,       snake_case	,       deterministic=not train						)

									if mid_block_additional_residual is not None:
											sample += mid_block_additional_residual

									# 5. up
									for up_block in self.up_blocks:
											_UpperCAmelCase							 =		down_block_res_samples[-(self.layers_per_block + 1) :]
											_UpperCAmelCase							 =		down_block_res_samples[: -(self.layers_per_block + 1)]
											if isinstance(snake_case	,       snake_case						):
													_UpperCAmelCase							 =		up_block(
													    snake_case	,       temb=snake_case	,       encoder_hidden_states=snake_case	,       res_hidden_states_tuple=snake_case	,       deterministic=not train	,       )
											else:
													_UpperCAmelCase							 =		up_block(snake_case	,       temb=snake_case	,       res_hidden_states_tuple=snake_case	,       deterministic=not train						)

        # 6. post-process
									_UpperCAmelCase							 =		self.conv_norm_out(snake_case						)
									_UpperCAmelCase							 =		nn.silu(snake_case						)
									_UpperCAmelCase							 =		self.conv_out(snake_case						)
									_UpperCAmelCase							 =		jnp.transpose(snake_case	,       (0, 3, 1, 2)						)

									if not return_dict:
											return (sample,)

									return FlaxUNetaDConditionOutput(sample=snake_case						)







"""simple docstring"""





import inspect
import unittest
from math import floor

from transformers import CvtConfig
from transformers.file_utils import cached_property, is_torch_available, is_vision_available
from transformers.testing_utils import require_torch, require_vision, slow, torch_device

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 CvtForImageClassification, CvtModel
					from transformers.models.cvt.modeling_cvt import CVT_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
					from PIL import Image

					from transformers import AutoImageProcessor




class 			lowercase__			(      A    ):
							'''simple docstring'''







							def        lowerCamelCase_   (  self						)			->							int:
									_UpperCAmelCase							 =		self.config_class(**self.inputs_dict						)
									self.parent.assertTrue(hasattr(snake_case	,       'embed_dim'						)						)
									self.parent.assertTrue(hasattr(snake_case	,       'num_heads'						)						)







class 			lowercase__			:
							'''simple docstring'''







							def __init__(  self	,       snake_case	,       snake_case=13	,       snake_case=64	,       snake_case=3	,       snake_case=[16, 48, 96]	,       snake_case=[1, 3, 6]	,       snake_case=[1, 2, 10]	,       snake_case=[7, 3, 3]	,       snake_case=[4, 2, 2]	,       snake_case=[2, 1, 1]	,       snake_case=[2, 2, 2]	,       snake_case=[False, False, True]	,       snake_case=[0.0, 0.0, 0.0]	,       snake_case=0.02	,       snake_case=1E-12	,       snake_case=True	,       snake_case=True	,       snake_case=2	,       )			->							Tuple:
									_UpperCAmelCase							 =		parent
									_UpperCAmelCase							 =		batch_size
									_UpperCAmelCase							 =		image_size
									_UpperCAmelCase							 =		patch_sizes
									_UpperCAmelCase							 =		patch_stride
									_UpperCAmelCase							 =		patch_padding
									_UpperCAmelCase							 =		is_training
									_UpperCAmelCase							 =		use_labels
									_UpperCAmelCase							 =		num_labels
									_UpperCAmelCase							 =		num_channels
									_UpperCAmelCase							 =		embed_dim
									_UpperCAmelCase							 =		num_heads
									_UpperCAmelCase							 =		stride_kv
									_UpperCAmelCase							 =		depth
									_UpperCAmelCase							 =		cls_token
									_UpperCAmelCase							 =		attention_drop_rate
									_UpperCAmelCase							 =		initializer_range
									_UpperCAmelCase							 =		layer_norm_eps






							def        lowerCamelCase_   (  self						)			->							Dict:
									_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.num_labels						)

									_UpperCAmelCase							 =		self.get_config()
									return config, pixel_values, labels






							def        lowerCamelCase_   (  self						)			->							List[str]:
									return CvtConfig(
									    image_size=self.image_size	,       num_labels=self.num_labels	,       num_channels=self.num_channels	,       embed_dim=self.embed_dim	,       num_heads=self.num_heads	,       patch_sizes=self.patch_sizes	,       patch_padding=self.patch_padding	,       patch_stride=self.patch_stride	,       stride_kv=self.stride_kv	,       depth=self.depth	,       cls_token=self.cls_token	,       attention_drop_rate=self.attention_drop_rate	,       initializer_range=self.initializer_range	,       )






							def        lowerCamelCase_   (  self	,       snake_case	,       snake_case	,       snake_case						)			->							Optional[int]:
									_UpperCAmelCase							 =		CvtModel(config=snake_case						)
									model.to(snake_case						)
									model.eval()
									_UpperCAmelCase							 =		model(snake_case						)
									_UpperCAmelCase							 =		(self.image_size, self.image_size)
									_UpperCAmelCase ,					_UpperCAmelCase							 =		image_size[0], image_size[1]
									for i in range(len(self.depth						)						):
											_UpperCAmelCase							 =		floor(((height + 2 * self.patch_padding[i] - self.patch_sizes[i]) / self.patch_stride[i]) + 1						)
											_UpperCAmelCase							 =		floor(((width + 2 * self.patch_padding[i] - self.patch_sizes[i]) / self.patch_stride[i]) + 1						)
									self.parent.assertEqual(result.last_hidden_state.shape	,       (self.batch_size, self.embed_dim[-1], height, width)						)






							def        lowerCamelCase_   (  self	,       snake_case	,       snake_case	,       snake_case						)			->							Optional[Any]:
									_UpperCAmelCase							 =		self.num_labels
									_UpperCAmelCase							 =		CvtForImageClassification(snake_case						)
									model.to(snake_case						)
									model.eval()
									_UpperCAmelCase							 =		model(snake_case	,       labels=snake_case						)
									self.parent.assertEqual(result.logits.shape	,       (self.batch_size, self.num_labels)						)






							def        lowerCamelCase_   (  self						)			->							Union[str, Any]:
									_UpperCAmelCase							 =		self.prepare_config_and_inputs()
									_UpperCAmelCase ,					_UpperCAmelCase ,					_UpperCAmelCase							 =		config_and_inputs
									_UpperCAmelCase							 =		{'pixel_values': pixel_values}
									return config, inputs_dict







@require_torch
class 			lowercase__			(      A,					A,					unittest.TestCase    ):
							'''simple docstring'''







							_UpperCAmelCase      				=					(CvtModel, CvtForImageClassification) if is_torch_available() else ()
							_UpperCAmelCase      				=					(
							    {'''feature-extraction''': CvtModel, '''image-classification''': CvtForImageClassification}
							    if is_torch_available()
							    else {}
							)

							_UpperCAmelCase      				=					False
							_UpperCAmelCase      				=					False
							_UpperCAmelCase      				=					False
							_UpperCAmelCase      				=					False
							_UpperCAmelCase      				=					False






							def        lowerCamelCase_   (  self						)			->							Union[str, Any]:
									_UpperCAmelCase							 =		CvtModelTester(self						)
									_UpperCAmelCase							 =		ConfigTester(self	,       config_class=snake_case	,       has_text_modality=snake_case	,       hidden_size=37						)






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






							@unittest.skip(reason='Cvt does not output attentions'						)
							def        lowerCamelCase_   (  self						)			->							str:
									pass






							@unittest.skip(reason='Cvt does not use inputs_embeds'						)
							def        lowerCamelCase_   (  self						)			->							int:
									pass






							@unittest.skip(reason='Cvt does not support input and output embeddings'						)
							def        lowerCamelCase_   (  self						)			->							Union[str, Any]:
									pass






							def        lowerCamelCase_   (  self						)			->							Any:
									_UpperCAmelCase ,					_UpperCAmelCase							 =		self.model_tester.prepare_config_and_inputs_for_common()

									for model_class in self.all_model_classes:
											_UpperCAmelCase							 =		model_class(snake_case						)
											_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]	,       snake_case						)






							def        lowerCamelCase_   (  self						)			->							Optional[int]:
									_UpperCAmelCase							 =		self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_model(*snake_case						)






							def        lowerCamelCase_   (  self						)			->							Optional[int]:
									def check_hidden_states_output(snake_case	,       snake_case	,       snake_case						):
											_UpperCAmelCase							 =		model_class(snake_case						)
											model.to(snake_case						)
											model.eval()

											with torch.no_grad():
													_UpperCAmelCase							 =		model(**self._prepare_for_class(snake_case	,       snake_case						)						)

											_UpperCAmelCase							 =		outputs.hidden_states

											_UpperCAmelCase							 =		len(self.model_tester.depth						)
											self.assertEqual(len(snake_case						)	,       snake_case						)

											# verify the first hidden states (first block)
											self.assertListEqual(
											    list(hidden_states[0].shape[-3:]						)	,       [
											        self.model_tester.embed_dim[0],
											        self.model_tester.image_size // 4,
											        self.model_tester.image_size // 4,
											    ]	,       )

									_UpperCAmelCase ,					_UpperCAmelCase							 =		self.model_tester.prepare_config_and_inputs_for_common()

									for model_class in self.all_model_classes:
											_UpperCAmelCase							 =		True
											check_hidden_states_output(snake_case	,       snake_case	,       snake_case						)

											# check that output_hidden_states also work using config
											del inputs_dict["output_hidden_states"]
											_UpperCAmelCase							 =		True

											check_hidden_states_output(snake_case	,       snake_case	,       snake_case						)






							def        lowerCamelCase_   (  self						)			->							Any:
									_UpperCAmelCase							 =		self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_for_image_classification(*snake_case						)






							@unittest.skip('Will be fixed soon by reducing the size of the model used for common tests.'						)
							def        lowerCamelCase_   (  self						)			->							Dict:
									pass






							@slow
							def        lowerCamelCase_   (  self						)			->							Dict:
									for model_name in CVT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
											_UpperCAmelCase							 =		CvtModel.from_pretrained(snake_case						)
											self.assertIsNotNone(snake_case						)







def 	UpperCAmelCase     (		):

		'''simple docstring'''


		_UpperCAmelCase							 =		Image.open('./tests/fixtures/tests_samples/COCO/000000039769.png'       )
		return image






@require_torch
@require_vision
class 			lowercase__			(      unittest.TestCase    ):
							'''simple docstring'''







							@cached_property
							def        lowerCamelCase_   (  self						)			->							List[Any]:
									return AutoImageProcessor.from_pretrained(CVT_PRETRAINED_MODEL_ARCHIVE_LIST[0]						)






							@slow
							def        lowerCamelCase_   (  self						)			->							Dict:
									_UpperCAmelCase							 =		CvtForImageClassification.from_pretrained(CVT_PRETRAINED_MODEL_ARCHIVE_LIST[0]						).to(snake_case						)

									_UpperCAmelCase							 =		self.default_image_processor
									_UpperCAmelCase							 =		prepare_img()
									_UpperCAmelCase							 =		image_processor(images=snake_case	,       return_tensors='pt'						).to(snake_case						)

									# forward pass
									with torch.no_grad():
											_UpperCAmelCase							 =		model(**snake_case						)

									# verify the logits
									_UpperCAmelCase							 =		torch.Size((1, 1000)						)
									self.assertEqual(outputs.logits.shape	,       snake_case						)

									_UpperCAmelCase							 =		torch.tensor([0.9285, 0.9015, -0.3150]						).to(snake_case						)

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


