

"""simple docstring"""






import flax.linen as nn
import jax.numpy as jnp

from .attention_flax import FlaxTransformeraDModel
from .resnet_flax import FlaxDownsampleaD, FlaxResnetBlockaD, FlaxUpsampleaD







class       snake_case__  ( nn.Module     ):
			SCREAMING_SNAKE_CASE__     				=						42
			SCREAMING_SNAKE_CASE__     				=						42
			SCREAMING_SNAKE_CASE__     				=						0.0
			SCREAMING_SNAKE_CASE__     				=						1
			SCREAMING_SNAKE_CASE__     				=						1
			SCREAMING_SNAKE_CASE__     				=						True
			SCREAMING_SNAKE_CASE__     				=						False
			SCREAMING_SNAKE_CASE__     				=						False
			SCREAMING_SNAKE_CASE__     				=						False
			SCREAMING_SNAKE_CASE__     				=						jnp.floataa







			def 			__lowerCAmelCase				(			self			: List[str]  ):





								'''simple docstring'''



								UpperCAmelCase			:   Tuple      =				[]
								UpperCAmelCase			:   Optional[Any]      =				[]

								for i in range(self.num_layers  ):
													UpperCAmelCase			:   List[str]      =				self.in_channels if i == 0 else self.out_channels

													UpperCAmelCase			:   List[Any]      =				FlaxResnetBlockaD(
													    in_channels=lowercase			,  out_channels=self.out_channels			,  dropout_prob=self.dropout			,  dtype=self.dtype			,  )
													resnets.append(lowercase  )

													UpperCAmelCase			:   Any      =				FlaxTransformeraDModel(
													    in_channels=self.out_channels			,  n_heads=self.num_attention_heads			,  d_head=self.out_channels // self.num_attention_heads			,  depth=1			,  use_linear_projection=self.use_linear_projection			,  only_cross_attention=self.only_cross_attention			,  use_memory_efficient_attention=self.use_memory_efficient_attention			,  dtype=self.dtype			,  )
													attentions.append(lowercase  )

								UpperCAmelCase			:   str      =				resnets
								UpperCAmelCase			:   List[str]      =				attentions

								if self.add_downsample:
													UpperCAmelCase			:   str      =				FlaxDownsampleaD(self.out_channels			,  dtype=self.dtype  )







			def __call__(			self			: Dict			,  lowercase			: List[str]			,  lowercase			: Dict			,  lowercase			: Tuple			,  lowercase			: Optional[int]=True  ):





								'''simple docstring'''



								UpperCAmelCase			:   Tuple      =				()

								for resnet, attn in zip(self.resnets			,  self.attentions  ):
													UpperCAmelCase			:   Dict      =				resnet(lowercase			,  lowercase			,  deterministic=lowercase  )
													UpperCAmelCase			:   Optional[Any]      =				attn(lowercase			,  lowercase			,  deterministic=lowercase  )
													output_states += (hidden_states,)

								if self.add_downsample:
													UpperCAmelCase			:   Any      =				self.downsamplers_a(lowercase  )
													output_states += (hidden_states,)

								return hidden_states, output_states







class       snake_case__  ( nn.Module     ):
			SCREAMING_SNAKE_CASE__     				=						42
			SCREAMING_SNAKE_CASE__     				=						42
			SCREAMING_SNAKE_CASE__     				=						0.0
			SCREAMING_SNAKE_CASE__     				=						1
			SCREAMING_SNAKE_CASE__     				=						True
			SCREAMING_SNAKE_CASE__     				=						jnp.floataa







			def 			__lowerCAmelCase				(			self			: Optional[Any]  ):





								'''simple docstring'''



								UpperCAmelCase			:   Optional[int]      =				[]

								for i in range(self.num_layers  ):
													UpperCAmelCase			:   List[str]      =				self.in_channels if i == 0 else self.out_channels

													UpperCAmelCase			:   int      =				FlaxResnetBlockaD(
													    in_channels=lowercase			,  out_channels=self.out_channels			,  dropout_prob=self.dropout			,  dtype=self.dtype			,  )
													resnets.append(lowercase  )
								UpperCAmelCase			:   int      =				resnets

								if self.add_downsample:
													UpperCAmelCase			:   int      =				FlaxDownsampleaD(self.out_channels			,  dtype=self.dtype  )







			def __call__(			self			: Tuple			,  lowercase			: Dict			,  lowercase			: Dict			,  lowercase			: Union[str, Any]=True  ):





								'''simple docstring'''



								UpperCAmelCase			:   Tuple      =				()

								for resnet in self.resnets:
													UpperCAmelCase			:   Optional[Any]      =				resnet(lowercase			,  lowercase			,  deterministic=lowercase  )
													output_states += (hidden_states,)

								if self.add_downsample:
													UpperCAmelCase			:   Optional[Any]      =				self.downsamplers_a(lowercase  )
													output_states += (hidden_states,)

								return hidden_states, output_states







class       snake_case__  ( nn.Module     ):
			SCREAMING_SNAKE_CASE__     				=						42
			SCREAMING_SNAKE_CASE__     				=						42
			SCREAMING_SNAKE_CASE__     				=						42
			SCREAMING_SNAKE_CASE__     				=						0.0
			SCREAMING_SNAKE_CASE__     				=						1
			SCREAMING_SNAKE_CASE__     				=						1
			SCREAMING_SNAKE_CASE__     				=						True
			SCREAMING_SNAKE_CASE__     				=						False
			SCREAMING_SNAKE_CASE__     				=						False
			SCREAMING_SNAKE_CASE__     				=						False
			SCREAMING_SNAKE_CASE__     				=						jnp.floataa







			def 			__lowerCAmelCase				(			self			: List[str]  ):





								'''simple docstring'''



								UpperCAmelCase			:   Any      =				[]
								UpperCAmelCase			:   int      =				[]

								for i in range(self.num_layers  ):
													UpperCAmelCase			:   List[Any]      =				self.in_channels if (i == self.num_layers - 1) else self.out_channels
													UpperCAmelCase			:   Dict      =				self.prev_output_channel if i == 0 else self.out_channels

													UpperCAmelCase			:   int      =				FlaxResnetBlockaD(
													    in_channels=resnet_in_channels + res_skip_channels			,  out_channels=self.out_channels			,  dropout_prob=self.dropout			,  dtype=self.dtype			,  )
													resnets.append(lowercase  )

													UpperCAmelCase			:   List[str]      =				FlaxTransformeraDModel(
													    in_channels=self.out_channels			,  n_heads=self.num_attention_heads			,  d_head=self.out_channels // self.num_attention_heads			,  depth=1			,  use_linear_projection=self.use_linear_projection			,  only_cross_attention=self.only_cross_attention			,  use_memory_efficient_attention=self.use_memory_efficient_attention			,  dtype=self.dtype			,  )
													attentions.append(lowercase  )

								UpperCAmelCase			:   int      =				resnets
								UpperCAmelCase			:   Optional[int]      =				attentions

								if self.add_upsample:
													UpperCAmelCase			:   List[str]      =				FlaxUpsampleaD(self.out_channels			,  dtype=self.dtype  )







			def __call__(			self			: Dict			,  lowercase			: Dict			,  lowercase			: Optional[int]			,  lowercase			: List[Any]			,  lowercase			: Dict			,  lowercase			: Union[str, Any]=True  ):





								'''simple docstring'''



								for resnet, attn in zip(self.resnets			,  self.attentions  ):
													# pop res hidden states
													UpperCAmelCase			:   List[Any]      =				res_hidden_states_tuple[-1]
													UpperCAmelCase			:   int      =				res_hidden_states_tuple[:-1]
													UpperCAmelCase			:   Tuple      =				jnp.concatenate((hidden_states, res_hidden_states)			,  axis=-1  )

													UpperCAmelCase			:   Dict      =				resnet(lowercase			,  lowercase			,  deterministic=lowercase  )
													UpperCAmelCase			:   Tuple      =				attn(lowercase			,  lowercase			,  deterministic=lowercase  )

								if self.add_upsample:
													UpperCAmelCase			:   List[str]      =				self.upsamplers_a(lowercase  )

								return hidden_states







class       snake_case__  ( nn.Module     ):
			SCREAMING_SNAKE_CASE__     				=						42
			SCREAMING_SNAKE_CASE__     				=						42
			SCREAMING_SNAKE_CASE__     				=						42
			SCREAMING_SNAKE_CASE__     				=						0.0
			SCREAMING_SNAKE_CASE__     				=						1
			SCREAMING_SNAKE_CASE__     				=						True
			SCREAMING_SNAKE_CASE__     				=						jnp.floataa







			def 			__lowerCAmelCase				(			self			: List[str]  ):





								'''simple docstring'''



								UpperCAmelCase			:   List[str]      =				[]

								for i in range(self.num_layers  ):
													UpperCAmelCase			:   List[Any]      =				self.in_channels if (i == self.num_layers - 1) else self.out_channels
													UpperCAmelCase			:   List[Any]      =				self.prev_output_channel if i == 0 else self.out_channels

													UpperCAmelCase			:   Optional[int]      =				FlaxResnetBlockaD(
													    in_channels=resnet_in_channels + res_skip_channels			,  out_channels=self.out_channels			,  dropout_prob=self.dropout			,  dtype=self.dtype			,  )
													resnets.append(lowercase  )

								UpperCAmelCase			:   Dict      =				resnets

								if self.add_upsample:
													UpperCAmelCase			:   List[str]      =				FlaxUpsampleaD(self.out_channels			,  dtype=self.dtype  )







			def __call__(			self			: List[Any]			,  lowercase			: Optional[Any]			,  lowercase			: Tuple			,  lowercase			: Optional[Any]			,  lowercase			: int=True  ):





								'''simple docstring'''



								for resnet in self.resnets:
													# pop res hidden states
													UpperCAmelCase			:   Union[str, Any]      =				res_hidden_states_tuple[-1]
													UpperCAmelCase			:   List[Any]      =				res_hidden_states_tuple[:-1]
													UpperCAmelCase			:   Any      =				jnp.concatenate((hidden_states, res_hidden_states)			,  axis=-1  )

													UpperCAmelCase			:   Any      =				resnet(lowercase			,  lowercase			,  deterministic=lowercase  )

								if self.add_upsample:
													UpperCAmelCase			:   List[str]      =				self.upsamplers_a(lowercase  )

								return hidden_states







class       snake_case__  ( nn.Module     ):
			SCREAMING_SNAKE_CASE__     				=						42
			SCREAMING_SNAKE_CASE__     				=						0.0
			SCREAMING_SNAKE_CASE__     				=						1
			SCREAMING_SNAKE_CASE__     				=						1
			SCREAMING_SNAKE_CASE__     				=						False
			SCREAMING_SNAKE_CASE__     				=						False
			SCREAMING_SNAKE_CASE__     				=						jnp.floataa







			def 			__lowerCAmelCase				(			self			: Union[str, Any]  ):





								'''simple docstring'''



								UpperCAmelCase			:   Optional[int]      =				[
								    FlaxResnetBlockaD(
								        in_channels=self.in_channels			,  out_channels=self.in_channels			,  dropout_prob=self.dropout			,  dtype=self.dtype			,  )
								]

								UpperCAmelCase			:   Dict      =				[]

								for _ in range(self.num_layers  ):
													UpperCAmelCase			:   List[Any]      =				FlaxTransformeraDModel(
													    in_channels=self.in_channels			,  n_heads=self.num_attention_heads			,  d_head=self.in_channels // self.num_attention_heads			,  depth=1			,  use_linear_projection=self.use_linear_projection			,  use_memory_efficient_attention=self.use_memory_efficient_attention			,  dtype=self.dtype			,  )
													attentions.append(lowercase  )

													UpperCAmelCase			:   List[str]      =				FlaxResnetBlockaD(
													    in_channels=self.in_channels			,  out_channels=self.in_channels			,  dropout_prob=self.dropout			,  dtype=self.dtype			,  )
													resnets.append(lowercase  )

								UpperCAmelCase			:   List[str]      =				resnets
								UpperCAmelCase			:   Optional[Any]      =				attentions







			def __call__(			self			: int			,  lowercase			: List[str]			,  lowercase			: List[Any]			,  lowercase			: Tuple			,  lowercase			: Optional[Any]=True  ):





								'''simple docstring'''



								UpperCAmelCase			:   Optional[Any]      =				self.resnets[0](lowercase			,  lowercase  )
								for attn, resnet in zip(self.attentions			,  self.resnets[1:]  ):
													UpperCAmelCase			:   str      =				attn(lowercase			,  lowercase			,  deterministic=lowercase  )
													UpperCAmelCase			:   List[str]      =				resnet(lowercase			,  lowercase			,  deterministic=lowercase  )

								return hidden_states





"""simple docstring"""






snake_case_   :   str        =       [
    """DownloadConfig""",
    """DownloadManager""",
    """DownloadMode""",
    """StreamingDownloadManager""",
]

from .download_config import DownloadConfig
from .download_manager import DownloadManager, DownloadMode
from .streaming_download_manager import StreamingDownloadManager


