


"""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 _lowerCAmelCase							( snake_case_      ):
   __UpperCAmelCase					:			jnp.ndarray







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







   def 					lowerCamelCase						(      self		,      UpperCamelCase__		)      ->						FrozenDict:

       '''simple docstring'''




       snake_case       : List[str]    =       (1, self.in_channels, self.sample_size, self.sample_size)
       snake_case       : List[Any]    =       jnp.zeros(UpperCamelCase__		,      dtype=jnp.floataa		)
       snake_case       : Optional[int]    =       jnp.ones((1,)		,      dtype=jnp.intaa		)
       snake_case       : Union[str, Any]    =       jnp.zeros((1, 1, self.cross_attention_dim)		,      dtype=jnp.floataa		)

       snake_case   ,snake_case       : List[Any]    =       jax.random.split(UpperCamelCase__		)
       snake_case       : int    =       {"params": params_rng, "dropout": dropout_rng}

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







   def 					lowerCamelCase						(      self		)      ->						List[str]:

       '''simple docstring'''




       snake_case       : Any    =       self.block_out_channels
       snake_case       : Optional[int]    =       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.
       snake_case       : List[Any]    =       self.num_attention_heads or self.attention_head_dim

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

       # time
       snake_case       : int    =       FlaxTimesteps(
           block_out_channels[0]		,      flip_sin_to_cos=self.flip_sin_to_cos		,      freq_shift=self.config.freq_shift		)
       snake_case       : Optional[Any]    =       FlaxTimestepEmbedding(UpperCamelCase__		,      dtype=self.dtype		)

       snake_case       : Optional[int]    =       self.only_cross_attention
       if isinstance(UpperCamelCase__		,      UpperCamelCase__		):
           snake_case       : List[Any]    =       (only_cross_attention,) * len(self.down_block_types		)

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

       # down
       snake_case       : Optional[Any]    =       []
       snake_case       : int    =       block_out_channels[0]
       for i, down_block_type in enumerate(self.down_block_types		):
           snake_case       : Optional[int]    =       output_channel
           snake_case       : List[str]    =       block_out_channels[i]
           snake_case       : List[Any]    =       i == len(UpperCamelCase__		) - 1

           if down_block_type == "CrossAttnDownBlock2D":
               snake_case       : Optional[int]    =       FlaxCrossAttnDownBlockaD(
                   in_channels=UpperCamelCase__		,      out_channels=UpperCamelCase__		,      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:
               snake_case       : Any    =       FlaxDownBlockaD(
                   in_channels=UpperCamelCase__		,      out_channels=UpperCamelCase__		,      dropout=self.dropout		,      num_layers=self.layers_per_block		,      add_downsample=not is_final_block		,      dtype=self.dtype		,      )

           down_blocks.append(UpperCamelCase__		)
       snake_case       : List[Any]    =       down_blocks

       # mid
       snake_case       : Any    =       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
       snake_case       : Any    =       []
       snake_case       : Union[str, Any]    =       list(reversed(UpperCamelCase__		)		)
       snake_case       : int    =       list(reversed(UpperCamelCase__		)		)
       snake_case       : str    =       list(reversed(UpperCamelCase__		)		)
       snake_case       : Optional[int]    =       reversed_block_out_channels[0]
       for i, up_block_type in enumerate(self.up_block_types		):
           snake_case       : Any    =       output_channel
           snake_case       : int    =       reversed_block_out_channels[i]
           snake_case       : Tuple    =       reversed_block_out_channels[min(i + 1		,      len(UpperCamelCase__		) - 1		)]

           snake_case       : Tuple    =       i == len(UpperCamelCase__		) - 1

           if up_block_type == "CrossAttnUpBlock2D":
               snake_case       : Optional[int]    =       FlaxCrossAttnUpBlockaD(
                   in_channels=UpperCamelCase__		,      out_channels=UpperCamelCase__		,      prev_output_channel=UpperCamelCase__		,      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:
               snake_case       : str    =       FlaxUpBlockaD(
                   in_channels=UpperCamelCase__		,      out_channels=UpperCamelCase__		,      prev_output_channel=UpperCamelCase__		,      num_layers=self.layers_per_block + 1		,      add_upsample=not is_final_block		,      dropout=self.dropout		,      dtype=self.dtype		,      )

           up_blocks.append(UpperCamelCase__		)
           snake_case       : List[str]    =       output_channel
       snake_case       : Any    =       up_blocks

       # out
       snake_case       : str    =       nn.GroupNorm(num_groups=32		,      epsilon=1e-5		)
       snake_case       : Optional[int]    =       nn.Conv(
           self.out_channels		,      kernel_size=(3, 3)		,      strides=(1, 1)		,      padding=((1, 1), (1, 1))		,      dtype=self.dtype		,      )







   def __call__(      self		,      UpperCamelCase__		,      UpperCamelCase__		,      UpperCamelCase__		,      UpperCamelCase__=None		,      UpperCamelCase__=None		,      UpperCamelCase__ = True		,      UpperCamelCase__ = False		,      )      ->						Union[FlaxUNetaDConditionOutput, Tuple]:

       '''simple docstring'''




       if not isinstance(UpperCamelCase__		,      jnp.ndarray		):
           snake_case       : Any    =       jnp.array([timesteps]		,      dtype=jnp.intaa		)
       elif isinstance(UpperCamelCase__		,      jnp.ndarray		) and len(timesteps.shape		) == 0:
           snake_case       : List[str]    =       timesteps.astype(dtype=jnp.floataa		)
           snake_case       : List[str]    =       jnp.expand_dims(UpperCamelCase__		,      0		)

       snake_case       : Any    =       self.time_proj(UpperCamelCase__		)
       snake_case       : Optional[Any]    =       self.time_embedding(UpperCamelCase__		)

       # 2. pre-process
       snake_case       : int    =       jnp.transpose(UpperCamelCase__		,      (0, 2, 3, 1)		)
       snake_case       : List[str]    =       self.conv_in(UpperCamelCase__		)

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

       if down_block_additional_residuals is not None:
           snake_case       : Optional[int]    =       ()

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

           snake_case       : Any    =       new_down_block_res_samples

       # 4. mid
       snake_case       : List[str]    =       self.mid_block(UpperCamelCase__		,      UpperCamelCase__		,      UpperCamelCase__		,      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:
           snake_case       : int    =       down_block_res_samples[-(self.layers_per_block + 1) :]
           snake_case       : Optional[int]    =       down_block_res_samples[: -(self.layers_per_block + 1)]
           if isinstance(UpperCamelCase__		,      UpperCamelCase__		):
               snake_case       : List[str]    =       up_block(
                   UpperCamelCase__		,      temb=UpperCamelCase__		,      encoder_hidden_states=UpperCamelCase__		,      res_hidden_states_tuple=UpperCamelCase__		,      deterministic=not train		,      )
           else:
               snake_case       : Dict    =       up_block(UpperCamelCase__		,      temb=UpperCamelCase__		,      res_hidden_states_tuple=UpperCamelCase__		,      deterministic=not train		)

        # 6. post-process
       snake_case       : Union[str, Any]    =       self.conv_norm_out(UpperCamelCase__		)
       snake_case       : Union[str, Any]    =       nn.silu(UpperCamelCase__		)
       snake_case       : List[str]    =       self.conv_out(UpperCamelCase__		)
       snake_case       : Any    =       jnp.transpose(UpperCamelCase__		,      (0, 3, 1, 2)		)

       if not return_dict:
           return (sample,)

       return FlaxUNetaDConditionOutput(sample=UpperCamelCase__		)









"""simple docstring"""



from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available


__snake_case       =       {
    """configuration_time_series_transformer""": [
        """TIME_SERIES_TRANSFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP""",
        """TimeSeriesTransformerConfig""",
    ],
}

try:
  if not is_torch_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  __snake_case       =       [
      """TIME_SERIES_TRANSFORMER_PRETRAINED_MODEL_ARCHIVE_LIST""",
      """TimeSeriesTransformerForPrediction""",
      """TimeSeriesTransformerModel""",
      """TimeSeriesTransformerPreTrainedModel""",
  ]


if TYPE_CHECKING:
  from .configuration_time_series_transformer import (
      TIME_SERIES_TRANSFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP,
      TimeSeriesTransformerConfig,
  )

  try:
    if not is_torch_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .modeling_time_series_transformer import (
        TIME_SERIES_TRANSFORMER_PRETRAINED_MODEL_ARCHIVE_LIST,
        TimeSeriesTransformerForPrediction,
        TimeSeriesTransformerModel,
        TimeSeriesTransformerPreTrainedModel,
    )

else:
  import sys

  __snake_case       =       _LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)





