



import argparse
import json

import gdown
import numpy as np
import torch
from huggingface_hub import hf_hub_download

from transformers import (
    VideoMAEConfig,
    VideoMAEForPreTraining,
    VideoMAEForVideoClassification,
    VideoMAEImageProcessor,
)




def 	__UpperCamelCase	(  _lowerCAmelCase		)       ->  Any:



    """simple docstring"""






    A :       Optional[int]          =   VideoMAEConfig()

    set_architecture_configs(_lowerCAmelCase      ,     _lowerCAmelCase		)

    if "finetuned" not in model_name:
        A :       List[str]          =   False

    if "finetuned" in model_name:
        A :       Union[str, Any]          =   """huggingface/label-files"""
        if "kinetics" in model_name:
            A :       Tuple          =   400
            A :       List[Any]          =   """kinetics400-id2label.json"""
        elif "ssv2" in model_name:
            A :       Dict          =   174
            A :       Optional[Any]          =   """something-something-v2-id2label.json"""
        else:
            raise ValueError("""Model name should either contain 'kinetics' or 'ssv2' in case it's fine-tuned."""		)
        A :       Union[str, Any]          =   json.load(open(hf_hub_download(_lowerCAmelCase      ,     _lowerCAmelCase      ,     repo_type="""dataset"""		)      ,     """r"""		)		)
        A :       Any          =   {int(_lowerCAmelCase		): v for k, v in idalabel.items()}
        A :       Union[str, Any]          =   idalabel
        A :       Tuple          =   {v: k for k, v in idalabel.items()}

    return config




def 	__UpperCamelCase	(  _lowerCAmelCase      ,     _lowerCAmelCase		)       ->  Any:



    """simple docstring"""






    if "small" in model_name:
        A :       Optional[int]          =   384
        A :       int          =   1536
        A :       int          =   12
        A :       List[str]          =   16
        A :       Tuple          =   12
        A :       Optional[Any]          =   3
        A :       List[str]          =   192
        A :       List[str]          =   768
    elif "large" in model_name:
        A :       Dict          =   1024
        A :       List[str]          =   4096
        A :       List[str]          =   24
        A :       Tuple          =   16
        A :       List[Any]          =   12
        A :       Optional[int]          =   8
        A :       Dict          =   512
        A :       Tuple          =   2048
    elif "huge" in model_name:
        A :       Tuple          =   1280
        A :       Tuple          =   5120
        A :       int          =   32
        A :       List[Any]          =   16
        A :       Any          =   12
        A :       int          =   8
        A :       List[str]          =   640
        A :       List[str]          =   2560
    elif "base" not in model_name:
        raise ValueError("""Model name should include either \"small\", \"base\", \"large\", or \"huge\""""		)




def 	__UpperCamelCase	(  _lowerCAmelCase		)       ->  Optional[Any]:



    """simple docstring"""






    if "encoder." in name:
        A :       Any          =   name.replace("""encoder."""      ,     """"""		)
    if "cls_token" in name:
        A :       int          =   name.replace("""cls_token"""      ,     """videomae.embeddings.cls_token"""		)
    if "decoder_pos_embed" in name:
        A :       Dict          =   name.replace("""decoder_pos_embed"""      ,     """decoder.decoder_pos_embed"""		)
    if "pos_embed" in name and "decoder" not in name:
        A :       Optional[int]          =   name.replace("""pos_embed"""      ,     """videomae.embeddings.position_embeddings"""		)
    if "patch_embed.proj" in name:
        A :       Tuple          =   name.replace("""patch_embed.proj"""      ,     """videomae.embeddings.patch_embeddings.projection"""		)
    if "patch_embed.norm" in name:
        A :       str          =   name.replace("""patch_embed.norm"""      ,     """videomae.embeddings.norm"""		)
    if "decoder.blocks" in name:
        A :       Tuple          =   name.replace("""decoder.blocks"""      ,     """decoder.decoder_layers"""		)
    if "blocks" in name:
        A :       Union[str, Any]          =   name.replace("""blocks"""      ,     """videomae.encoder.layer"""		)
    if "attn.proj" in name:
        A :       Any          =   name.replace("""attn.proj"""      ,     """attention.output.dense"""		)
    if "attn" in name and "bias" not in name:
        A :       List[str]          =   name.replace("""attn"""      ,     """attention.self"""		)
    if "attn" in name:
        A :       Optional[Any]          =   name.replace("""attn"""      ,     """attention.attention"""		)
    if "norm1" in name:
        A :       Any          =   name.replace("""norm1"""      ,     """layernorm_before"""		)
    if "norm2" in name:
        A :       List[str]          =   name.replace("""norm2"""      ,     """layernorm_after"""		)
    if "mlp.fc1" in name:
        A :       Dict          =   name.replace("""mlp.fc1"""      ,     """intermediate.dense"""		)
    if "mlp.fc2" in name:
        A :       List[str]          =   name.replace("""mlp.fc2"""      ,     """output.dense"""		)
    if "decoder_embed" in name:
        A :       str          =   name.replace("""decoder_embed"""      ,     """decoder.decoder_embed"""		)
    if "decoder_norm" in name:
        A :       Optional[Any]          =   name.replace("""decoder_norm"""      ,     """decoder.decoder_norm"""		)
    if "decoder_pred" in name:
        A :       Optional[int]          =   name.replace("""decoder_pred"""      ,     """decoder.decoder_pred"""		)
    if "norm.weight" in name and "decoder" not in name and "fc" not in name:
        A :       Dict          =   name.replace("""norm.weight"""      ,     """videomae.layernorm.weight"""		)
    if "norm.bias" in name and "decoder" not in name and "fc" not in name:
        A :       Optional[int]          =   name.replace("""norm.bias"""      ,     """videomae.layernorm.bias"""		)
    if "head" in name and "decoder" not in name:
        A :       List[str]          =   name.replace("""head"""      ,     """classifier"""		)

    return name




def 	__UpperCamelCase	(  _lowerCAmelCase      ,     _lowerCAmelCase		)       ->  List[Any]:



    """simple docstring"""






    for key in orig_state_dict.copy().keys():
        A :       List[Any]          =   orig_state_dict.pop(_lowerCAmelCase		)

        if key.startswith("""encoder."""		):
            A :       int          =   key.replace("""encoder."""      ,     """"""		)

        if "qkv" in key:
            A :       List[Any]          =   key.split("""."""		)
            if key.startswith("""decoder.blocks"""		):
                A :       Dict          =   config.decoder_hidden_size
                A :       Dict          =   int(key_split[2]		)
                A :       Dict          =   """decoder.decoder_layers."""
                if "weight" in key:
                    A :       Optional[int]          =   val[:dim, :]
                    A :       List[Any]          =   val[dim : dim * 2, :]
                    A :       List[str]          =   val[-dim:, :]
            else:
                A :       Optional[Any]          =   config.hidden_size
                A :       Optional[int]          =   int(key_split[1]		)
                A :       Optional[Any]          =   """videomae.encoder.layer."""
                if "weight" in key:
                    A :       Dict          =   val[:dim, :]
                    A :       Any          =   val[dim : dim * 2, :]
                    A :       Any          =   val[-dim:, :]
        else:
            A :       Dict          =   val

    return orig_state_dict




def 	__UpperCamelCase	(  )       ->  Optional[Any]:



    """simple docstring"""






    A :       Any          =   hf_hub_download(
        repo_id="""hf-internal-testing/spaghetti-video"""      ,     filename="""eating_spaghetti.npy"""      ,     repo_type="""dataset"""		)
    A :       Optional[Any]          =   np.load(_lowerCAmelCase		)
    return list(_lowerCAmelCase		)




def 	__UpperCamelCase	(  _lowerCAmelCase      ,     _lowerCAmelCase      ,     _lowerCAmelCase      ,     _lowerCAmelCase		)       ->  Dict:



    """simple docstring"""






    A :       str          =   get_videomae_config(_lowerCAmelCase		)

    if "finetuned" in model_name:
        A :       Any          =   VideoMAEForVideoClassification(_lowerCAmelCase		)
    else:
        A :       Optional[Any]          =   VideoMAEForPreTraining(_lowerCAmelCase		)

    # download original checkpoint, hosted on Google Drive
    A :       Any          =   """pytorch_model.bin"""
    gdown.cached_download(_lowerCAmelCase      ,     _lowerCAmelCase      ,     quiet=_lowerCAmelCase		)
    A :       Dict          =   torch.load(_lowerCAmelCase      ,     map_location="""cpu"""		)
    if "model" in files:
        A :       List[str]          =   files["""model"""]
    else:
        A :       Optional[int]          =   files["""module"""]
    A :       str          =   convert_state_dict(_lowerCAmelCase      ,     _lowerCAmelCase		)

    model.load_state_dict(_lowerCAmelCase		)
    model.eval()

    # verify model on basic input
    A :       str          =   VideoMAEImageProcessor(image_mean=[0.5, 0.5, 0.5]      ,     image_std=[0.5, 0.5, 0.5]		)
    A :       List[Any]          =   prepare_video()
    A :       Optional[int]          =   image_processor(_lowerCAmelCase      ,     return_tensors="""pt"""		)

    if "finetuned" not in model_name:
        A :       Optional[int]          =   hf_hub_download(repo_id="""hf-internal-testing/bool-masked-pos"""      ,     filename="""bool_masked_pos.pt"""		)
        A :       Optional[Any]          =   torch.load(_lowerCAmelCase		)

    A :       int          =   model(**_lowerCAmelCase		)
    A :       Tuple          =   outputs.logits

    A :       Union[str, Any]          =   [
        """videomae-small-finetuned-kinetics""",
        """videomae-small-finetuned-ssv2""",
        # Kinetics-400 checkpoints (short = pretrained only for 800 epochs instead of 1600)
        """videomae-base-short""",
        """videomae-base-short-finetuned-kinetics""",
        """videomae-base""",
        """videomae-base-finetuned-kinetics""",
        """videomae-large""",
        """videomae-large-finetuned-kinetics""",
        """videomae-huge-finetuned-kinetics""",
        # Something-Something-v2 checkpoints (short = pretrained only for 800 epochs instead of 2400)
        """videomae-base-short-ssv2""",
        """videomae-base-short-finetuned-ssv2""",
        """videomae-base-ssv2""",
        """videomae-base-finetuned-ssv2""",
    ]

    # NOTE: logits were tested with image_mean and image_std equal to [0.5, 0.5, 0.5] and [0.5, 0.5, 0.5]
    if model_name == "videomae-small-finetuned-kinetics":
        A :       List[Any]          =   torch.Size([1, 400]		)
        A :       Optional[Any]          =   torch.tensor([-0.9_291, -0.4_061, -0.9_307]		)
    elif model_name == "videomae-small-finetuned-ssv2":
        A :       Union[str, Any]          =   torch.Size([1, 174]		)
        A :       List[str]          =   torch.tensor([0.2_671, -0.4_689, -0.8_235]		)
    elif model_name == "videomae-base":
        A :       List[Any]          =   torch.Size([1, 1408, 1536]		)
        A :       str          =   torch.tensor([[0.7_739, 0.7_968, 0.7_089], [0.6_701, 0.7_487, 0.6_209], [0.4_287, 0.5_158, 0.4_773]]		)
    elif model_name == "videomae-base-short":
        A :       List[str]          =   torch.Size([1, 1408, 1536]		)
        A :       List[str]          =   torch.tensor([[0.7_994, 0.9_612, 0.8_508], [0.7_401, 0.8_958, 0.8_302], [0.5_862, 0.7_468, 0.7_325]]		)
        # we verified the loss both for normalized and unnormalized targets for this one
        A :       Optional[Any]          =   torch.tensor([0.5_142]		) if config.norm_pix_loss else torch.tensor([0.6_469]		)
    elif model_name == "videomae-large":
        A :       Tuple          =   torch.Size([1, 1408, 1536]		)
        A :       Tuple          =   torch.tensor([[0.7_149, 0.7_997, 0.6_966], [0.6_768, 0.7_869, 0.6_948], [0.5_139, 0.6_221, 0.5_605]]		)
    elif model_name == "videomae-large-finetuned-kinetics":
        A :       Any          =   torch.Size([1, 400]		)
        A :       List[Any]          =   torch.tensor([0.0_771, 0.0_011, -0.3_625]		)
    elif model_name == "videomae-huge-finetuned-kinetics":
        A :       List[str]          =   torch.Size([1, 400]		)
        A :       Union[str, Any]          =   torch.tensor([0.2_433, 0.1_632, -0.4_894]		)
    elif model_name == "videomae-base-short-finetuned-kinetics":
        A :       Dict          =   torch.Size([1, 400]		)
        A :       Optional[int]          =   torch.tensor([0.6_588, 0.0_990, -0.2_493]		)
    elif model_name == "videomae-base-finetuned-kinetics":
        A :       int          =   torch.Size([1, 400]		)
        A :       Union[str, Any]          =   torch.tensor([0.3_669, -0.0_688, -0.2_421]		)
    elif model_name == "videomae-base-short-ssv2":
        A :       List[str]          =   torch.Size([1, 1408, 1536]		)
        A :       Dict          =   torch.tensor([[0.4_712, 0.5_296, 0.5_786], [0.2_278, 0.2_729, 0.4_026], [0.0_352, 0.0_730, 0.2_506]]		)
    elif model_name == "videomae-base-short-finetuned-ssv2":
        A :       str          =   torch.Size([1, 174]		)
        A :       Any          =   torch.tensor([-0.0_537, -0.1_539, -0.3_266]		)
    elif model_name == "videomae-base-ssv2":
        A :       Union[str, Any]          =   torch.Size([1, 1408, 1536]		)
        A :       List[str]          =   torch.tensor([[0.8_131, 0.8_727, 0.8_546], [0.7_366, 0.9_377, 0.8_870], [0.5_935, 0.8_874, 0.8_564]]		)
    elif model_name == "videomae-base-finetuned-ssv2":
        A :       Optional[Any]          =   torch.Size([1, 174]		)
        A :       List[Any]          =   torch.tensor([0.1_961, -0.8_337, -0.6_389]		)
    else:
        raise ValueError(f'''Model name not supported. Should be one of {model_names}'''		)

    # verify logits
    assert logits.shape == expected_shape
    if "finetuned" in model_name:
        assert torch.allclose(logits[0, :3]      ,     _lowerCAmelCase      ,     atol=1e-4		)
    else:
        print("""Logits:"""      ,     logits[0, :3, :3]		)
        assert torch.allclose(logits[0, :3, :3]      ,     _lowerCAmelCase      ,     atol=1e-4		)
    print("""Logits ok!"""		)

    # verify loss, if applicable
    if model_name == "videomae-base-short":
        A :       Tuple          =   outputs.loss
        assert torch.allclose(_lowerCAmelCase      ,     _lowerCAmelCase      ,     atol=1e-4		)
        print("""Loss ok!"""		)

    if pytorch_dump_folder_path is not None:
        print(f'''Saving model and image processor to {pytorch_dump_folder_path}'''		)
        image_processor.save_pretrained(_lowerCAmelCase		)
        model.save_pretrained(_lowerCAmelCase		)

    if push_to_hub:
        print("""Pushing to the hub..."""		)
        model.push_to_hub(_lowerCAmelCase      ,     organization="""nielsr"""		)


if __name__ == "__main__":
     SCREAMING_SNAKE_CASE_:int             =						argparse.ArgumentParser()
     # Required parameters
     parser.add_argument(
         """--checkpoint_url""",
         default="""https://drive.google.com/u/1/uc?id=1tEhLyskjb755TJ65ptsrafUG2llSwQE1&amp;export=download&amp;confirm=t&amp;uuid=aa3276eb-fb7e-482a-adec-dc7171df14c4""",
         type=str,
         help=(
             """URL of the original PyTorch checkpoint (on Google Drive) you'd like to convert. Should be a direct"""
             """ download link."""
         ),
     )
     parser.add_argument(
         """--pytorch_dump_folder_path""",
         default="""/Users/nielsrogge/Documents/VideoMAE/Test""",
         type=str,
         help="""Path to the output PyTorch model directory.""",
     )
     parser.add_argument("""--model_name""", default="""videomae-base""", type=str, help="""Name of the model.""")
     parser.add_argument(
         """--push_to_hub""", action="""store_true""", help="""Whether or not to push the converted model to the 🤗 hub."""
     )

     SCREAMING_SNAKE_CASE_:int             =						parser.parse_args()
     convert_videomae_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path, args.model_name, args.push_to_hub)







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 SCREAMING_SNAKE_CASE__     (		nn.Module    ):




    '''simple docstring'''






    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__()
        A :       Dict          =   only_cross_attention

        A :       Tuple          =   (num_embeds_ada_norm is not None) and norm_type == """ada_norm_zero"""
        A :       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:
            A :       Dict          =   AdaLayerNorm(lowerCamelCase__,		lowerCamelCase__		)
        elif self.use_ada_layer_norm_zero:
            A :       List[str]          =   AdaLayerNormZero(lowerCamelCase__,		lowerCamelCase__		)
        else:
            A :       Tuple          =   nn.LayerNorm(lowerCamelCase__,		elementwise_affine=lowerCamelCase__		)
        A :       Any          =   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.
            A :       int          =   (
                AdaLayerNorm(lowerCamelCase__,		lowerCamelCase__		)
                if self.use_ada_layer_norm
                else nn.LayerNorm(lowerCamelCase__,		elementwise_affine=lowerCamelCase__		)
            )
            A :       Dict          =   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:
            A :       Dict          =   None
            A :       Dict          =   None

        # 3. Feed-forward
        A :       Optional[Any]          =   nn.LayerNorm(lowerCamelCase__,		elementwise_affine=lowerCamelCase__		)
        A :       int          =   FeedForward(lowerCamelCase__,		dropout=lowerCamelCase__,		activation_fn=lowerCamelCase__,		final_dropout=lowerCamelCase__		)

        # let chunk size default to None
        A :       Optional[Any]          =   None
        A :       int          =   0






    def 			_lowerCAmelCase					(       self,		lowerCamelCase__,		lowerCamelCase__		):
        # Sets chunk feed-forward
        A :       List[str]          =   chunk_size
        A :       int          =   dim






    def 			_lowerCAmelCase					(       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:
            A :       Optional[int]          =   self.norma(lowerCamelCase__,		lowerCamelCase__		)
        elif self.use_ada_layer_norm_zero:
            A     ,							A     ,							A     ,							A     ,							A :       Tuple          =   self.norma(
                lowerCamelCase__,		lowerCamelCase__,		lowerCamelCase__,		hidden_dtype=hidden_states.dtype		)
        else:
            A :       Tuple          =   self.norma(lowerCamelCase__		)

        A :       Dict          =   cross_attention_kwargs if cross_attention_kwargs is not None else {}

        A :       str          =   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:
            A :       List[Any]          =   gate_msa.unsqueeze(1		) * attn_output
        A :       str          =   attn_output + hidden_states

        # 2. Cross-Attention
        if self.attna is not None:
            A :       Optional[int]          =   (
                self.norma(lowerCamelCase__,		lowerCamelCase__		) if self.use_ada_layer_norm else self.norma(lowerCamelCase__		)
            )

            A :       Optional[Any]          =   self.attna(
                lowerCamelCase__,		encoder_hidden_states=lowerCamelCase__,		attention_mask=lowerCamelCase__,		**lowerCamelCase__,		)
            A :       Dict          =   attn_output + hidden_states

        # 3. Feed-forward
        A :       str          =   self.norma(lowerCamelCase__		)

        if self.use_ada_layer_norm_zero:
            A :       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`.'''		)

            A :       Any          =   norm_hidden_states.shape[self._chunk_dim] // self._chunk_size
            A :       Optional[Any]          =   torch.cat(
                [self.ff(lowerCamelCase__		) for hid_slice in norm_hidden_states.chunk(lowerCamelCase__,		dim=self._chunk_dim		)],		dim=self._chunk_dim,		)
        else:
            A :       Dict          =   self.ff(lowerCamelCase__		)

        if self.use_ada_layer_norm_zero:
            A :       Optional[Any]          =   gate_mlp.unsqueeze(1		) * ff_output

        A :       Optional[int]          =   ff_output + hidden_states

        return hidden_states






class SCREAMING_SNAKE_CASE__     (		nn.Module    ):




    '''simple docstring'''






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

        if activation_fn == "gelu":
            A :       Dict          =   GELU(lowerCamelCase__,		lowerCamelCase__		)
        if activation_fn == "gelu-approximate":
            A :       Optional[int]          =   GELU(lowerCamelCase__,		lowerCamelCase__,		approximate="""tanh"""		)
        elif activation_fn == "geglu":
            A :       Any          =   GEGLU(lowerCamelCase__,		lowerCamelCase__		)
        elif activation_fn == "geglu-approximate":
            A :       Optional[Any]          =   ApproximateGELU(lowerCamelCase__,		lowerCamelCase__		)

        A :       Dict          =   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 			_lowerCAmelCase					(       self,		lowerCamelCase__		):
        for module in self.net:
            A :       int          =   module(lowerCamelCase__		)
        return hidden_states






class SCREAMING_SNAKE_CASE__     (		nn.Module    ):




    '''simple docstring'''






    def __init__(       self,		lowerCamelCase__,		lowerCamelCase__,		lowerCamelCase__ = "none"		):
        super().__init__()
        A :       Optional[int]          =   nn.Linear(lowerCamelCase__,		lowerCamelCase__		)
        A :       int          =   approximate






    def 			_lowerCAmelCase					(       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 			_lowerCAmelCase					(       self,		lowerCamelCase__		):
        A :       List[Any]          =   self.proj(lowerCamelCase__		)
        A :       Union[str, Any]          =   self.gelu(lowerCamelCase__		)
        return hidden_states






class SCREAMING_SNAKE_CASE__     (		nn.Module    ):




    '''simple docstring'''






    def __init__(       self,		lowerCamelCase__,		lowerCamelCase__		):
        super().__init__()
        A :       Union[str, Any]          =   nn.Linear(lowerCamelCase__,		dim_out * 2		)






    def 			_lowerCAmelCase					(       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 			_lowerCAmelCase					(       self,		lowerCamelCase__		):
        A     ,							A :       Any          =   self.proj(lowerCamelCase__		).chunk(2,		dim=-1		)
        return hidden_states * self.gelu(lowerCamelCase__		)






class SCREAMING_SNAKE_CASE__     (		nn.Module    ):




    '''simple docstring'''






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






    def 			_lowerCAmelCase					(       self,		lowerCamelCase__		):
        A :       List[Any]          =   self.proj(lowerCamelCase__		)
        return x * torch.sigmoid(1.702 * x		)






class SCREAMING_SNAKE_CASE__     (		nn.Module    ):




    '''simple docstring'''






    def __init__(       self,		lowerCamelCase__,		lowerCamelCase__		):
        super().__init__()
        A :       Dict          =   nn.Embedding(lowerCamelCase__,		lowerCamelCase__		)
        A :       Tuple          =   nn.SiLU()
        A :       Tuple          =   nn.Linear(lowerCamelCase__,		embedding_dim * 2		)
        A :       List[str]          =   nn.LayerNorm(lowerCamelCase__,		elementwise_affine=lowerCamelCase__		)






    def 			_lowerCAmelCase					(       self,		lowerCamelCase__,		lowerCamelCase__		):
        A :       int          =   self.linear(self.silu(self.emb(lowerCamelCase__		)		)		)
        A     ,							A :       Optional[int]          =   torch.chunk(lowerCamelCase__,		2		)
        A :       Tuple          =   self.norm(lowerCamelCase__		) * (1 + scale) + shift
        return x






class SCREAMING_SNAKE_CASE__     (		nn.Module    ):




    '''simple docstring'''






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

        A :       Union[str, Any]          =   CombinedTimestepLabelEmbeddings(lowerCamelCase__,		lowerCamelCase__		)

        A :       List[Any]          =   nn.SiLU()
        A :       str          =   nn.Linear(lowerCamelCase__,		6 * embedding_dim,		bias=lowerCamelCase__		)
        A :       List[Any]          =   nn.LayerNorm(lowerCamelCase__,		elementwise_affine=lowerCamelCase__,		eps=1e-6		)






    def 			_lowerCAmelCase					(       self,		lowerCamelCase__,		lowerCamelCase__,		lowerCamelCase__,		lowerCamelCase__=None		):
        A :       Tuple          =   self.linear(self.silu(self.emb(lowerCamelCase__,		lowerCamelCase__,		hidden_dtype=lowerCamelCase__		)		)		)
        A     ,							A     ,							A     ,							A     ,							A     ,							A :       Optional[Any]          =   emb.chunk(6,		dim=1		)
        A :       List[str]          =   self.norm(lowerCamelCase__		) * (1 + scale_msa[:, None]) + shift_msa[:, None]
        return x, gate_msa, shift_mlp, scale_mlp, gate_mlp






class SCREAMING_SNAKE_CASE__     (		nn.Module    ):




    '''simple docstring'''






    def __init__(       self,		lowerCamelCase__,		lowerCamelCase__,		lowerCamelCase__,		lowerCamelCase__ = None,		lowerCamelCase__ = 1e-5		):
        super().__init__()
        A :       int          =   num_groups
        A :       Dict          =   eps

        if act_fn is None:
            A :       Union[str, Any]          =   None
        else:
            A :       Optional[int]          =   get_activation(lowerCamelCase__		)

        A :       Dict          =   nn.Linear(lowerCamelCase__,		out_dim * 2		)






    def 			_lowerCAmelCase					(       self,		lowerCamelCase__,		lowerCamelCase__		):
        if self.act:
            A :       Dict          =   self.act(lowerCamelCase__		)
        A :       Optional[int]          =   self.linear(lowerCamelCase__		)
        A :       int          =   emb[:, :, None, None]
        A     ,							A :       Tuple          =   emb.chunk(2,		dim=1		)

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


