
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    lowerCAmelCase__     (	__lowercase			):
		a__       :						jnp.ndarray





@flax_register_to_config
class    lowerCAmelCase__     (	nn.Module  ,	__lowercase  ,	__lowercase			):
		a__       :						int        =		32
		a__       :						int        =		4
		a__       :						int        =		4
		a__       :						Tuple[str]        =		(
		    "CrossAttnDownBlock2D",
		    "CrossAttnDownBlock2D",
		    "CrossAttnDownBlock2D",
		    "DownBlock2D",
		)
		a__       :						Tuple[str]        =		("UpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D")
		a__       :						Union[bool, Tuple[bool]]        =		False
		a__       :						Tuple[int]        =		(320, 640, 1_280, 1_280)
		a__       :						int        =		2
		a__       :						Union[int, Tuple[int]]        =		8
		a__       :						Optional[Union[int, Tuple[int]]]        =		None
		a__       :						int        =		1_280
		a__       :						float        =		0.0
		a__       :						bool        =		False
		a__       :						jnp.dtype        =		jnp.floataa
		a__       :						bool        =		True
		a__       :						int        =		0
		a__       :						bool        =		False






		def  __A       (		self		:							str							,			SCREAMING_SNAKE_CASE__		:							jax.random.KeyArray  )     ->							FrozenDict:
						# init input tensors
						__lowerCamelCase    					=	(1, self.in_channels, self.sample_size, self.sample_size)
						__lowerCamelCase    					=	jnp.zeros(SCREAMING_SNAKE_CASE__							,			dtype=jnp.floataa  )
						__lowerCamelCase    					=	jnp.ones((1,)							,			dtype=jnp.intaa  )
						__lowerCamelCase    					=	jnp.zeros((1, 1, self.cross_attention_dim)							,			dtype=jnp.floataa  )

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

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






		def  __A       (		self		:							Union[str, Any]  )     ->							Any:
						__lowerCamelCase    					=	self.block_out_channels
						__lowerCamelCase    					=	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.
						__lowerCamelCase    					=	self.num_attention_heads or self.attention_head_dim

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

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

						__lowerCamelCase    					=	self.only_cross_attention
						if isinstance(SCREAMING_SNAKE_CASE__							,			SCREAMING_SNAKE_CASE__  ):
										__lowerCamelCase    					=	(only_cross_attention,) * len(self.down_block_types  )

						if isinstance(SCREAMING_SNAKE_CASE__							,			SCREAMING_SNAKE_CASE__  ):
										__lowerCamelCase    					=	(num_attention_heads,) * len(self.down_block_types  )

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

										if down_block_type == "CrossAttnDownBlock2D":
														__lowerCamelCase    					=	FlaxCrossAttnDownBlockaD(
														    in_channels=SCREAMING_SNAKE_CASE__							,			out_channels=SCREAMING_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:
														__lowerCamelCase    					=	FlaxDownBlockaD(
														    in_channels=SCREAMING_SNAKE_CASE__							,			out_channels=SCREAMING_SNAKE_CASE__							,			dropout=self.dropout							,			num_layers=self.layers_per_block							,			add_downsample=not is_final_block							,			dtype=self.dtype							,			)

										down_blocks.append(SCREAMING_SNAKE_CASE__  )
						__lowerCamelCase    					=	down_blocks

						# mid
						__lowerCamelCase    					=	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
						__lowerCamelCase    					=	[]
						__lowerCamelCase    					=	list(reversed(SCREAMING_SNAKE_CASE__  )  )
						__lowerCamelCase    					=	list(reversed(SCREAMING_SNAKE_CASE__  )  )
						__lowerCamelCase    					=	list(reversed(SCREAMING_SNAKE_CASE__  )  )
						__lowerCamelCase    					=	reversed_block_out_channels[0]
						for i, up_block_type in enumerate(self.up_block_types  ):
										__lowerCamelCase    					=	output_channel
										__lowerCamelCase    					=	reversed_block_out_channels[i]
										__lowerCamelCase    					=	reversed_block_out_channels[min(i + 1							,			len(SCREAMING_SNAKE_CASE__  ) - 1  )]

										__lowerCamelCase    					=	i == len(SCREAMING_SNAKE_CASE__  ) - 1

										if up_block_type == "CrossAttnUpBlock2D":
														__lowerCamelCase    					=	FlaxCrossAttnUpBlockaD(
														    in_channels=SCREAMING_SNAKE_CASE__							,			out_channels=SCREAMING_SNAKE_CASE__							,			prev_output_channel=SCREAMING_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:
														__lowerCamelCase    					=	FlaxUpBlockaD(
														    in_channels=SCREAMING_SNAKE_CASE__							,			out_channels=SCREAMING_SNAKE_CASE__							,			prev_output_channel=SCREAMING_SNAKE_CASE__							,			num_layers=self.layers_per_block + 1							,			add_upsample=not is_final_block							,			dropout=self.dropout							,			dtype=self.dtype							,			)

										up_blocks.append(SCREAMING_SNAKE_CASE__  )
										__lowerCamelCase    					=	output_channel
						__lowerCamelCase    					=	up_blocks

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







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

						__lowerCamelCase    					=	self.time_proj(SCREAMING_SNAKE_CASE__  )
						__lowerCamelCase    					=	self.time_embedding(SCREAMING_SNAKE_CASE__  )

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

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

						if down_block_additional_residuals is not None:
										__lowerCamelCase    					=	()

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

										__lowerCamelCase    					=	new_down_block_res_samples

						# 4. mid
						__lowerCamelCase    					=	self.mid_block(SCREAMING_SNAKE_CASE__							,			SCREAMING_SNAKE_CASE__							,			SCREAMING_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:
										__lowerCamelCase    					=	down_block_res_samples[-(self.layers_per_block + 1) :]
										__lowerCamelCase    					=	down_block_res_samples[: -(self.layers_per_block + 1)]
										if isinstance(SCREAMING_SNAKE_CASE__							,			SCREAMING_SNAKE_CASE__  ):
														__lowerCamelCase    					=	up_block(
														    SCREAMING_SNAKE_CASE__							,			temb=SCREAMING_SNAKE_CASE__							,			encoder_hidden_states=SCREAMING_SNAKE_CASE__							,			res_hidden_states_tuple=SCREAMING_SNAKE_CASE__							,			deterministic=not train							,			)
										else:
														__lowerCamelCase    					=	up_block(SCREAMING_SNAKE_CASE__							,			temb=SCREAMING_SNAKE_CASE__							,			res_hidden_states_tuple=SCREAMING_SNAKE_CASE__							,			deterministic=not train  )

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

						if not return_dict:
										return (sample,)

						return FlaxUNetaDConditionOutput(sample=SCREAMING_SNAKE_CASE__  )





import copy
from typing import Dict, List, Optional

from ...configuration_utils import PretrainedConfig
from ...utils import logging
from ..auto import CONFIG_MAPPING


SCREAMING_SNAKE_CASE__   :			Dict												=    {
    "facebook/mask2former-swin-small-coco-instance": (
        "https://huggingface.co/facebook/mask2former-swin-small-coco-instance/blob/main/config.json"
    )
    # See all Mask2Former models at https://huggingface.co/models?filter=mask2former
}

SCREAMING_SNAKE_CASE__   :			int												=    logging.get_logger(__name__)

class    lowerCAmelCase__     (	__lowercase			):
		a__       :						Any        =		"""mask2former"""
		a__       :						Dict        =		["""swin"""]
		a__       :						Any        =		{"""hidden_size""": """hidden_dim"""}






		def __init__(		self		:							Any							,			SCREAMING_SNAKE_CASE__		:							Optional[Dict] = None							,			SCREAMING_SNAKE_CASE__		:							int = 2_56							,			SCREAMING_SNAKE_CASE__		:							int = 2_56							,			SCREAMING_SNAKE_CASE__		:							int = 2_56							,			SCREAMING_SNAKE_CASE__		:							int = 10_24							,			SCREAMING_SNAKE_CASE__		:							str = "relu"							,			SCREAMING_SNAKE_CASE__		:							int = 6							,			SCREAMING_SNAKE_CASE__		:							int = 10							,			SCREAMING_SNAKE_CASE__		:							int = 8							,			SCREAMING_SNAKE_CASE__		:							float = 0.0							,			SCREAMING_SNAKE_CASE__		:							int = 20_48							,			SCREAMING_SNAKE_CASE__		:							bool = False							,			SCREAMING_SNAKE_CASE__		:							bool = False							,			SCREAMING_SNAKE_CASE__		:							int = 4							,			SCREAMING_SNAKE_CASE__		:							int = 2_55							,			SCREAMING_SNAKE_CASE__		:							int = 1_00							,			SCREAMING_SNAKE_CASE__		:							float = 0.1							,			SCREAMING_SNAKE_CASE__		:							float = 2.0							,			SCREAMING_SNAKE_CASE__		:							float = 5.0							,			SCREAMING_SNAKE_CASE__		:							float = 5.0							,			SCREAMING_SNAKE_CASE__		:							int = 1_25_44							,			SCREAMING_SNAKE_CASE__		:							float = 3.0							,			SCREAMING_SNAKE_CASE__		:							float = 0.75							,			SCREAMING_SNAKE_CASE__		:							float = 0.02							,			SCREAMING_SNAKE_CASE__		:							float = 1.0							,			SCREAMING_SNAKE_CASE__		:							bool = True							,			SCREAMING_SNAKE_CASE__		:							List[int] = [4, 8, 16, 32]							,			SCREAMING_SNAKE_CASE__		:							bool = None							,			**SCREAMING_SNAKE_CASE__		:							Tuple							,			)     ->							str:
						if backbone_config is None:
										logger.info('''`backbone_config` is `None`. Initializing the config with the default `Swin` backbone.'''  )
										__lowerCamelCase    					=	CONFIG_MAPPING['''swin'''](
										    image_size=2_24							,			in_channels=3							,			patch_size=4							,			embed_dim=96							,			depths=[2, 2, 18, 2]							,			num_heads=[3, 6, 12, 24]							,			window_size=7							,			drop_path_rate=0.3							,			use_absolute_embeddings=SCREAMING_SNAKE_CASE__							,			out_features=['''stage1''', '''stage2''', '''stage3''', '''stage4''']							,			)

						if isinstance(SCREAMING_SNAKE_CASE__							,			SCREAMING_SNAKE_CASE__  ):
										__lowerCamelCase    					=	backbone_config.pop('''model_type'''  )
										__lowerCamelCase    					=	CONFIG_MAPPING[backbone_model_type]
										__lowerCamelCase    					=	config_class.from_dict(SCREAMING_SNAKE_CASE__  )

						# verify that the backbone is supported
						if backbone_config.model_type not in self.backbones_supported:
										logger.warning_once(
										    f'''Backbone {backbone_config.model_type} is not a supported model and may not be compatible with Mask2Former. '''
										    f'''Supported model types: {','.join(self.backbones_supported  )}'''  )

						__lowerCamelCase    					=	backbone_config
						__lowerCamelCase    					=	feature_size
						__lowerCamelCase    					=	mask_feature_size
						__lowerCamelCase    					=	hidden_dim
						__lowerCamelCase    					=	encoder_feedforward_dim
						__lowerCamelCase    					=	activation_function
						__lowerCamelCase    					=	encoder_layers
						__lowerCamelCase    					=	decoder_layers
						__lowerCamelCase    					=	num_attention_heads
						__lowerCamelCase    					=	dropout
						__lowerCamelCase    					=	dim_feedforward
						__lowerCamelCase    					=	pre_norm
						__lowerCamelCase    					=	enforce_input_projection
						__lowerCamelCase    					=	common_stride
						__lowerCamelCase    					=	ignore_value
						__lowerCamelCase    					=	num_queries
						__lowerCamelCase    					=	no_object_weight
						__lowerCamelCase    					=	class_weight
						__lowerCamelCase    					=	mask_weight
						__lowerCamelCase    					=	dice_weight
						__lowerCamelCase    					=	train_num_points
						__lowerCamelCase    					=	oversample_ratio
						__lowerCamelCase    					=	importance_sample_ratio
						__lowerCamelCase    					=	init_std
						__lowerCamelCase    					=	init_xavier_std
						__lowerCamelCase    					=	use_auxiliary_loss
						__lowerCamelCase    					=	feature_strides
						__lowerCamelCase    					=	output_auxiliary_logits
						__lowerCamelCase    					=	decoder_layers

						super().__init__(**SCREAMING_SNAKE_CASE__  )






		@classmethod
		def  __A       (		cls		:							Any							,			SCREAMING_SNAKE_CASE__		:							PretrainedConfig							,			**SCREAMING_SNAKE_CASE__		:							Optional[Any]  )     ->							List[Any]:
						return cls(
						    backbone_config=SCREAMING_SNAKE_CASE__							,			**SCREAMING_SNAKE_CASE__							,			)







		def  __A       (		self		:							Any  )     ->							Dict[str, any]:
						__lowerCamelCase    					=	copy.deepcopy(self.__dict__  )
						__lowerCamelCase    					=	self.backbone_config.to_dict()
						__lowerCamelCase    					=	self.__class__.model_type
						return output



