


'''simple docstring'''



from typing import Any, Dict, Optional

import torch
import torch.nn.functional as F
from torch import nn

from ..utils import maybe_allow_in_graph
from .activations import get_activation
from .attention_processor import Attention
from .embeddings import CombinedTimestepLabelEmbeddings


@maybe_allow_in_graph
class 		_snake_case					(     nn.Module	):




					def __init__(      self							,							_lowerCamelCase							,							_lowerCamelCase							,							_lowerCamelCase							,							_lowerCamelCase=0.0							,							_lowerCamelCase = None							,							_lowerCamelCase = "geglu"							,							_lowerCamelCase = None							,							_lowerCamelCase = False							,							_lowerCamelCase = False							,							_lowerCamelCase = False							,							_lowerCamelCase = False							,							_lowerCamelCase = True							,							_lowerCamelCase = "layer_norm"							,							_lowerCamelCase = False							,							):
								super().__init__()
								UpperCAmelCase__  :					Optional[Any]        =     only_cross_attention

								UpperCAmelCase__  :					Dict        =     (num_embeds_ada_norm is not None) and norm_type == """ada_norm_zero"""
								UpperCAmelCase__  :					Union[str, Any]        =     (num_embeds_ada_norm is not None) and norm_type == """ada_norm"""

								if norm_type in ("ada_norm", "ada_norm_zero") and num_embeds_ada_norm is None:
											raise ValueError(
											    f'''`norm_type` is set to {norm_type}, but `num_embeds_ada_norm` is not defined. Please make sure to'''
											    f''' define `num_embeds_ada_norm` if setting `norm_type` to {norm_type}.''')

								# Define 3 blocks. Each block has its own normalization layer.
								# 1. Self-Attn
								if self.use_ada_layer_norm:
											UpperCAmelCase__  :					Dict        =     AdaLayerNorm(_lowerCamelCase							,							_lowerCamelCase)
								elif self.use_ada_layer_norm_zero:
											UpperCAmelCase__  :					Tuple        =     AdaLayerNormZero(_lowerCamelCase							,							_lowerCamelCase)
								else:
											UpperCAmelCase__  :					List[str]        =     nn.LayerNorm(_lowerCamelCase							,							elementwise_affine=_lowerCamelCase)
								UpperCAmelCase__  :					Tuple        =     Attention(
								    query_dim=_lowerCamelCase							,							heads=_lowerCamelCase							,							dim_head=_lowerCamelCase							,							dropout=_lowerCamelCase							,							bias=_lowerCamelCase							,							cross_attention_dim=cross_attention_dim if only_cross_attention else None							,							upcast_attention=_lowerCamelCase							,							)

								# 2. Cross-Attn
								if cross_attention_dim is not None or double_self_attention:
											# We currently only use AdaLayerNormZero for self attention where there will only be one attention block.
											# I.e. the number of returned modulation chunks from AdaLayerZero would not make sense if returned during
											# the second cross attention block.
											UpperCAmelCase__  :					Any        =     (
											    AdaLayerNorm(_lowerCamelCase							,							_lowerCamelCase)
											    if self.use_ada_layer_norm
											    else nn.LayerNorm(_lowerCamelCase							,							elementwise_affine=_lowerCamelCase)
											)
											UpperCAmelCase__  :					List[str]        =     Attention(
											    query_dim=_lowerCamelCase							,							cross_attention_dim=cross_attention_dim if not double_self_attention else None							,							heads=_lowerCamelCase							,							dim_head=_lowerCamelCase							,							dropout=_lowerCamelCase							,							bias=_lowerCamelCase							,							upcast_attention=_lowerCamelCase							,							)  # is self-attn if encoder_hidden_states is none
								else:
											UpperCAmelCase__  :					int        =     None
											UpperCAmelCase__  :					Tuple        =     None

								# 3. Feed-forward
								UpperCAmelCase__  :					List[Any]        =     nn.LayerNorm(_lowerCamelCase							,							elementwise_affine=_lowerCamelCase)
								UpperCAmelCase__  :					int        =     FeedForward(_lowerCamelCase							,							dropout=_lowerCamelCase							,							activation_fn=_lowerCamelCase							,							final_dropout=_lowerCamelCase)

								# let chunk size default to None
								UpperCAmelCase__  :					Tuple        =     None
								UpperCAmelCase__  :					List[Any]        =     0




					def 				snake_case__						(      self							,							_lowerCamelCase							,							_lowerCamelCase):
								# Sets chunk feed-forward
								UpperCAmelCase__  :					Optional[Any]        =     chunk_size
								UpperCAmelCase__  :					Optional[int]        =     dim




					def 				snake_case__						(      self							,							_lowerCamelCase							,							_lowerCamelCase = None							,							_lowerCamelCase = None							,							_lowerCamelCase = None							,							_lowerCamelCase = None							,							_lowerCamelCase = None							,							_lowerCamelCase = None							,							):
								# Notice that normalization is always applied before the real computation in the following blocks.
								# 1. Self-Attention
								if self.use_ada_layer_norm:
											UpperCAmelCase__  :					Any        =     self.norma(_lowerCamelCase							,							_lowerCamelCase)
								elif self.use_ada_layer_norm_zero:
											UpperCAmelCase__   ,       UpperCAmelCase__   ,       UpperCAmelCase__   ,       UpperCAmelCase__   ,       UpperCAmelCase__  :					int        =     self.norma(
											    _lowerCamelCase							,							_lowerCamelCase							,							_lowerCamelCase							,							hidden_dtype=hidden_states.dtype)
								else:
											UpperCAmelCase__  :					Tuple        =     self.norma(_lowerCamelCase)

								UpperCAmelCase__  :					Any        =     cross_attention_kwargs if cross_attention_kwargs is not None else {}

								UpperCAmelCase__  :					Dict        =     self.attna(
								    _lowerCamelCase							,							encoder_hidden_states=encoder_hidden_states if self.only_cross_attention else None							,							attention_mask=_lowerCamelCase							,							**_lowerCamelCase							,							)
								if self.use_ada_layer_norm_zero:
											UpperCAmelCase__  :					Any        =     gate_msa.unsqueeze(1) * attn_output
								UpperCAmelCase__  :					int        =     attn_output + hidden_states

								# 2. Cross-Attention
								if self.attna is not None:
											UpperCAmelCase__  :					str        =     (
											    self.norma(_lowerCamelCase							,							_lowerCamelCase) if self.use_ada_layer_norm else self.norma(_lowerCamelCase)
											)

											UpperCAmelCase__  :					Union[str, Any]        =     self.attna(
											    _lowerCamelCase							,							encoder_hidden_states=_lowerCamelCase							,							attention_mask=_lowerCamelCase							,							**_lowerCamelCase							,							)
											UpperCAmelCase__  :					int        =     attn_output + hidden_states

								# 3. Feed-forward
								UpperCAmelCase__  :					int        =     self.norma(_lowerCamelCase)

								if self.use_ada_layer_norm_zero:
											UpperCAmelCase__  :					Tuple        =     norm_hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None]

								if self._chunk_size is not None:
											# "feed_forward_chunk_size" can be used to save memory
											if norm_hidden_states.shape[self._chunk_dim] % self._chunk_size != 0:
														raise ValueError(
														    f'''`hidden_states` dimension to be chunked: {norm_hidden_states.shape[self._chunk_dim]} has to be divisible by chunk size: {self._chunk_size}. Make sure to set an appropriate `chunk_size` when calling `unet.enable_forward_chunking`.''')

											UpperCAmelCase__  :					List[str]        =     norm_hidden_states.shape[self._chunk_dim] // self._chunk_size
											UpperCAmelCase__  :					List[Any]        =     torch.cat(
											    [self.ff(_lowerCamelCase) for hid_slice in norm_hidden_states.chunk(_lowerCamelCase							,							dim=self._chunk_dim)]							,							dim=self._chunk_dim							,							)
								else:
											UpperCAmelCase__  :					Dict        =     self.ff(_lowerCamelCase)

								if self.use_ada_layer_norm_zero:
											UpperCAmelCase__  :					str        =     gate_mlp.unsqueeze(1) * ff_output

								UpperCAmelCase__  :					Optional[int]        =     ff_output + hidden_states

								return hidden_states


class 		_snake_case					(     nn.Module	):




					def __init__(      self							,							_lowerCamelCase							,							_lowerCamelCase = None							,							_lowerCamelCase = 4							,							_lowerCamelCase = 0.0							,							_lowerCamelCase = "geglu"							,							_lowerCamelCase = False							,							):
								super().__init__()
								UpperCAmelCase__  :					Union[str, Any]        =     int(dim * mult)
								UpperCAmelCase__  :					Optional[Any]        =     dim_out if dim_out is not None else dim

								if activation_fn == "gelu":
											UpperCAmelCase__  :					List[str]        =     GELU(_lowerCamelCase							,							_lowerCamelCase)
								if activation_fn == "gelu-approximate":
											UpperCAmelCase__  :					Tuple        =     GELU(_lowerCamelCase							,							_lowerCamelCase							,							approximate="""tanh""")
								elif activation_fn == "geglu":
											UpperCAmelCase__  :					str        =     GEGLU(_lowerCamelCase							,							_lowerCamelCase)
								elif activation_fn == "geglu-approximate":
											UpperCAmelCase__  :					str        =     ApproximateGELU(_lowerCamelCase							,							_lowerCamelCase)

								UpperCAmelCase__  :					str        =     nn.ModuleList([])
								# project in
								self.net.append(_lowerCamelCase)
								# project dropout
								self.net.append(nn.Dropout(_lowerCamelCase))
								# project out
								self.net.append(nn.Linear(_lowerCamelCase							,							_lowerCamelCase))
								# FF as used in Vision Transformer, MLP-Mixer, etc. have a final dropout
								if final_dropout:
											self.net.append(nn.Dropout(_lowerCamelCase))




					def 				snake_case__						(      self							,							_lowerCamelCase):
								for module in self.net:
											UpperCAmelCase__  :					Union[str, Any]        =     module(_lowerCamelCase)
								return hidden_states


class 		_snake_case					(     nn.Module	):




					def __init__(      self							,							_lowerCamelCase							,							_lowerCamelCase							,							_lowerCamelCase = "none"):
								super().__init__()
								UpperCAmelCase__  :					Union[str, Any]        =     nn.Linear(_lowerCamelCase							,							_lowerCamelCase)
								UpperCAmelCase__  :					Dict        =     approximate




					def 				snake_case__						(      self							,							_lowerCamelCase):
								if gate.device.type != "mps":
											return F.gelu(_lowerCamelCase							,							approximate=self.approximate)
								# mps: gelu is not implemented for float16
								return F.gelu(gate.to(dtype=torch.floataa)							,							approximate=self.approximate).to(dtype=gate.dtype)




					def 				snake_case__						(      self							,							_lowerCamelCase):
								UpperCAmelCase__  :					str        =     self.proj(_lowerCamelCase)
								UpperCAmelCase__  :					List[Any]        =     self.gelu(_lowerCamelCase)
								return hidden_states


class 		_snake_case					(     nn.Module	):




					def __init__(      self							,							_lowerCamelCase							,							_lowerCamelCase):
								super().__init__()
								UpperCAmelCase__  :					Dict        =     nn.Linear(_lowerCamelCase							,							dim_out * 2)




					def 				snake_case__						(      self							,							_lowerCamelCase):
								if gate.device.type != "mps":
											return F.gelu(_lowerCamelCase)
								# mps: gelu is not implemented for float16
								return F.gelu(gate.to(dtype=torch.floataa)).to(dtype=gate.dtype)




					def 				snake_case__						(      self							,							_lowerCamelCase):
								UpperCAmelCase__   ,       UpperCAmelCase__  :					List[str]        =     self.proj(_lowerCamelCase).chunk(2							,							dim=-1)
								return hidden_states * self.gelu(_lowerCamelCase)


class 		_snake_case					(     nn.Module	):




					def __init__(      self							,							_lowerCamelCase							,							_lowerCamelCase):
								super().__init__()
								UpperCAmelCase__  :					Union[str, Any]        =     nn.Linear(_lowerCamelCase							,							_lowerCamelCase)




					def 				snake_case__						(      self							,							_lowerCamelCase):
								UpperCAmelCase__  :					Dict        =     self.proj(_lowerCamelCase)
								return x * torch.sigmoid(1.702 * x)


class 		_snake_case					(     nn.Module	):




					def __init__(      self							,							_lowerCamelCase							,							_lowerCamelCase):
								super().__init__()
								UpperCAmelCase__  :					List[str]        =     nn.Embedding(_lowerCamelCase							,							_lowerCamelCase)
								UpperCAmelCase__  :					Optional[int]        =     nn.SiLU()
								UpperCAmelCase__  :					Union[str, Any]        =     nn.Linear(_lowerCamelCase							,							embedding_dim * 2)
								UpperCAmelCase__  :					List[Any]        =     nn.LayerNorm(_lowerCamelCase							,							elementwise_affine=_lowerCamelCase)




					def 				snake_case__						(      self							,							_lowerCamelCase							,							_lowerCamelCase):
								UpperCAmelCase__  :					Tuple        =     self.linear(self.silu(self.emb(_lowerCamelCase)))
								UpperCAmelCase__   ,       UpperCAmelCase__  :					Any        =     torch.chunk(_lowerCamelCase							,							2)
								UpperCAmelCase__  :					int        =     self.norm(_lowerCamelCase) * (1 + scale) + shift
								return x


class 		_snake_case					(     nn.Module	):




					def __init__(      self							,							_lowerCamelCase							,							_lowerCamelCase):
								super().__init__()

								UpperCAmelCase__  :					Dict        =     CombinedTimestepLabelEmbeddings(_lowerCamelCase							,							_lowerCamelCase)

								UpperCAmelCase__  :					Dict        =     nn.SiLU()
								UpperCAmelCase__  :					List[Any]        =     nn.Linear(_lowerCamelCase							,							6 * embedding_dim							,							bias=_lowerCamelCase)
								UpperCAmelCase__  :					Optional[int]        =     nn.LayerNorm(_lowerCamelCase							,							elementwise_affine=_lowerCamelCase							,							eps=1e-6)




					def 				snake_case__						(      self							,							_lowerCamelCase							,							_lowerCamelCase							,							_lowerCamelCase							,							_lowerCamelCase=None):
								UpperCAmelCase__  :					Tuple        =     self.linear(self.silu(self.emb(_lowerCamelCase							,							_lowerCamelCase							,							hidden_dtype=_lowerCamelCase)))
								UpperCAmelCase__   ,       UpperCAmelCase__   ,       UpperCAmelCase__   ,       UpperCAmelCase__   ,       UpperCAmelCase__   ,       UpperCAmelCase__  :					Union[str, Any]        =     emb.chunk(6							,							dim=1)
								UpperCAmelCase__  :					int        =     self.norm(_lowerCamelCase) * (1 + scale_msa[:, None]) + shift_msa[:, None]
								return x, gate_msa, shift_mlp, scale_mlp, gate_mlp


class 		_snake_case					(     nn.Module	):




					def __init__(      self							,							_lowerCamelCase							,							_lowerCamelCase							,							_lowerCamelCase							,							_lowerCamelCase = None							,							_lowerCamelCase = 1e-5):
								super().__init__()
								UpperCAmelCase__  :					Any        =     num_groups
								UpperCAmelCase__  :					Tuple        =     eps

								if act_fn is None:
											UpperCAmelCase__  :					Optional[int]        =     None
								else:
											UpperCAmelCase__  :					str        =     get_activation(_lowerCamelCase)

								UpperCAmelCase__  :					int        =     nn.Linear(_lowerCamelCase							,							out_dim * 2)




					def 				snake_case__						(      self							,							_lowerCamelCase							,							_lowerCamelCase):
								if self.act:
											UpperCAmelCase__  :					Union[str, Any]        =     self.act(_lowerCamelCase)
								UpperCAmelCase__  :					List[str]        =     self.linear(_lowerCamelCase)
								UpperCAmelCase__  :					int        =     emb[:, :, None, None]
								UpperCAmelCase__   ,       UpperCAmelCase__  :					Optional[Any]        =     emb.chunk(2							,							dim=1)

								UpperCAmelCase__  :					List[Any]        =     F.group_norm(_lowerCamelCase							,							self.num_groups							,							eps=self.eps)
								UpperCAmelCase__  :					Union[str, Any]        =     x * (1 + scale) + shift
								return x



'''simple docstring'''



from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_torch_available,
)


__A	={
    'configuration_encodec': [
        'ENCODEC_PRETRAINED_CONFIG_ARCHIVE_MAP',
        'EncodecConfig',
    ],
    'feature_extraction_encodec': ['EncodecFeatureExtractor'],
}

try:
		if not is_torch_available():
				raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
		pass
else:
		__A	=[
		    'ENCODEC_PRETRAINED_MODEL_ARCHIVE_LIST',
		    'EncodecModel',
		    'EncodecPreTrainedModel',
		]

if TYPE_CHECKING:
		from .configuration_encodec import (
		    ENCODEC_PRETRAINED_CONFIG_ARCHIVE_MAP,
		    EncodecConfig,
		)
		from .feature_extraction_encodec import EncodecFeatureExtractor

		try:
				if not is_torch_available():
						raise OptionalDependencyNotAvailable()
		except OptionalDependencyNotAvailable:
				pass
		else:
				from .modeling_encodec import (
				    ENCODEC_PRETRAINED_MODEL_ARCHIVE_LIST,
				    EncodecModel,
				    EncodecPreTrainedModel,
				)

else:
		import sys

		__A	=_LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)