

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 					__lowercase      (     _A			)     ->						Optional[Any]:
      SCREAMING_SNAKE_CASE :      Optional[int]       =   VideoMAEConfig()

      set_architecture_configs(_A						,    _A			)

      if "finetuned" not in model_name:
            SCREAMING_SNAKE_CASE :      Dict       =   False

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

      return config






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






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

      return name






def 					__lowercase      (     _A						,    _A			)     ->						str:
      for key in orig_state_dict.copy().keys():
            SCREAMING_SNAKE_CASE :      Optional[Any]       =   orig_state_dict.pop(_A			)

            if key.startswith("""encoder."""			):
                  SCREAMING_SNAKE_CASE :      Dict       =   key.replace("""encoder."""						,    """"""			)

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

      return orig_state_dict






def 					__lowercase      (     )     ->						Union[str, Any]:
      SCREAMING_SNAKE_CASE :      List[str]       =   hf_hub_download(
          repo_id="""hf-internal-testing/spaghetti-video"""						,    filename="""eating_spaghetti.npy"""						,    repo_type="""dataset"""			)
      SCREAMING_SNAKE_CASE :      Dict       =   np.load(_A			)
      return list(_A			)






def 					__lowercase      (     _A						,    _A						,    _A						,    _A			)     ->						List[str]:
      SCREAMING_SNAKE_CASE :      Union[str, Any]       =   get_videomae_config(_A			)

      if "finetuned" in model_name:
            SCREAMING_SNAKE_CASE :      str       =   VideoMAEForVideoClassification(_A			)
      else:
            SCREAMING_SNAKE_CASE :      Any       =   VideoMAEForPreTraining(_A			)

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

      model.load_state_dict(_A			)
      model.eval()

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

      if "finetuned" not in model_name:
            SCREAMING_SNAKE_CASE :      int       =   hf_hub_download(repo_id="""hf-internal-testing/bool-masked-pos"""						,    filename="""bool_masked_pos.pt"""			)
            SCREAMING_SNAKE_CASE :      List[str]       =   torch.load(_A			)

      SCREAMING_SNAKE_CASE :      str       =   model(**_A			)
      SCREAMING_SNAKE_CASE :      str       =   outputs.logits

      SCREAMING_SNAKE_CASE :      int       =   [
          """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":
            SCREAMING_SNAKE_CASE :      Union[str, Any]       =   torch.Size([1, 400]			)
            SCREAMING_SNAKE_CASE :      List[Any]       =   torch.tensor([-0.9291, -0.4061, -0.9307]			)
      elif model_name == "videomae-small-finetuned-ssv2":
            SCREAMING_SNAKE_CASE :      List[Any]       =   torch.Size([1, 174]			)
            SCREAMING_SNAKE_CASE :      Union[str, Any]       =   torch.tensor([0.2671, -0.4689, -0.8235]			)
      elif model_name == "videomae-base":
            SCREAMING_SNAKE_CASE :      List[str]       =   torch.Size([1, 1408, 1536]			)
            SCREAMING_SNAKE_CASE :      Dict       =   torch.tensor([[0.7739, 0.7968, 0.7089], [0.6701, 0.7487, 0.6209], [0.4287, 0.5158, 0.4773]]			)
      elif model_name == "videomae-base-short":
            SCREAMING_SNAKE_CASE :      Dict       =   torch.Size([1, 1408, 1536]			)
            SCREAMING_SNAKE_CASE :      List[Any]       =   torch.tensor([[0.7994, 0.9612, 0.8508], [0.7401, 0.8958, 0.8302], [0.5862, 0.7468, 0.7325]]			)
            # we verified the loss both for normalized and unnormalized targets for this one
            SCREAMING_SNAKE_CASE :      List[str]       =   torch.tensor([0.5142]			) if config.norm_pix_loss else torch.tensor([0.6469]			)
      elif model_name == "videomae-large":
            SCREAMING_SNAKE_CASE :      str       =   torch.Size([1, 1408, 1536]			)
            SCREAMING_SNAKE_CASE :      Optional[int]       =   torch.tensor([[0.7149, 0.7997, 0.6966], [0.6768, 0.7869, 0.6948], [0.5139, 0.6221, 0.5605]]			)
      elif model_name == "videomae-large-finetuned-kinetics":
            SCREAMING_SNAKE_CASE :      Any       =   torch.Size([1, 400]			)
            SCREAMING_SNAKE_CASE :      Optional[int]       =   torch.tensor([0.0771, 0.0011, -0.3625]			)
      elif model_name == "videomae-huge-finetuned-kinetics":
            SCREAMING_SNAKE_CASE :      Dict       =   torch.Size([1, 400]			)
            SCREAMING_SNAKE_CASE :      str       =   torch.tensor([0.2433, 0.1632, -0.4894]			)
      elif model_name == "videomae-base-short-finetuned-kinetics":
            SCREAMING_SNAKE_CASE :      Tuple       =   torch.Size([1, 400]			)
            SCREAMING_SNAKE_CASE :      Optional[Any]       =   torch.tensor([0.6588, 0.0990, -0.2493]			)
      elif model_name == "videomae-base-finetuned-kinetics":
            SCREAMING_SNAKE_CASE :      Optional[Any]       =   torch.Size([1, 400]			)
            SCREAMING_SNAKE_CASE :      List[str]       =   torch.tensor([0.3669, -0.0688, -0.2421]			)
      elif model_name == "videomae-base-short-ssv2":
            SCREAMING_SNAKE_CASE :      Any       =   torch.Size([1, 1408, 1536]			)
            SCREAMING_SNAKE_CASE :      Optional[Any]       =   torch.tensor([[0.4712, 0.5296, 0.5786], [0.2278, 0.2729, 0.4026], [0.0352, 0.0730, 0.2506]]			)
      elif model_name == "videomae-base-short-finetuned-ssv2":
            SCREAMING_SNAKE_CASE :      Any       =   torch.Size([1, 174]			)
            SCREAMING_SNAKE_CASE :      Optional[Any]       =   torch.tensor([-0.0537, -0.1539, -0.3266]			)
      elif model_name == "videomae-base-ssv2":
            SCREAMING_SNAKE_CASE :      Any       =   torch.Size([1, 1408, 1536]			)
            SCREAMING_SNAKE_CASE :      Optional[int]       =   torch.tensor([[0.8131, 0.8727, 0.8546], [0.7366, 0.9377, 0.8870], [0.5935, 0.8874, 0.8564]]			)
      elif model_name == "videomae-base-finetuned-ssv2":
            SCREAMING_SNAKE_CASE :      Dict       =   torch.Size([1, 174]			)
            SCREAMING_SNAKE_CASE :      Union[str, Any]       =   torch.tensor([0.1961, -0.8337, -0.6389]			)
      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]						,    _A						,    atol=1e-4			)
      else:
            print("""Logits:"""						,    logits[0, :3, :3]			)
            assert torch.allclose(logits[0, :3, :3]						,    _A						,    atol=1e-4			)
      print("""Logits ok!"""			)

      # verify loss, if applicable
      if model_name == "videomae-base-short":
            SCREAMING_SNAKE_CASE :      Optional[Any]       =   outputs.loss
            assert torch.allclose(_A						,    _A						,    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(_A			)
            model.save_pretrained(_A			)

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


if __name__ == "__main__":
     UpperCAmelCase__ :     Tuple				    =    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."""
     )

     UpperCAmelCase__ :     List[str]				    =    parser.parse_args()
     convert_videomae_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path, args.model_name, args.push_to_hub)







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

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







class        a__       (   nn.Module       ):







  """simple docstring"""




  UpperCAmelCase__      : int
  UpperCAmelCase__      : int
  UpperCAmelCase__      : float			       =0.0
  UpperCAmelCase__      : int			       =1
  UpperCAmelCase__      : int			       =1
  UpperCAmelCase__      : bool			       =True
  UpperCAmelCase__      : bool			       =False
  UpperCAmelCase__      : bool			       =False
  UpperCAmelCase__      : bool			       =False
  UpperCAmelCase__      : jnp.dtype			       =jnp.floataa


  def  _lowercase      (			self    : Tuple		)					->Optional[Any]:
        """simple docstring"""







        SCREAMING_SNAKE_CASE :      List[str]       =   []
        SCREAMING_SNAKE_CASE :      Any       =   []

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

              SCREAMING_SNAKE_CASE :      Union[str, Any]       =   FlaxResnetBlockaD(
                  in_channels=UpperCAmelCase__				,				out_channels=self.out_channels				,				dropout_prob=self.dropout				,				dtype=self.dtype				,				)
              resnets.append(UpperCAmelCase__		)

              SCREAMING_SNAKE_CASE :      List[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(UpperCAmelCase__		)

        SCREAMING_SNAKE_CASE :      Optional[Any]       =   resnets
        SCREAMING_SNAKE_CASE :      Optional[int]       =   attentions

        if self.add_downsample:
              SCREAMING_SNAKE_CASE :      Tuple       =   FlaxDownsampleaD(self.out_channels				,				dtype=self.dtype		)




  def __call__(			self    : int				,				UpperCAmelCase__    : List[str]				,				UpperCAmelCase__    : Union[str, Any]				,				UpperCAmelCase__    : Optional[Any]				,				UpperCAmelCase__    : List[str]=True		)					->str:
        """simple docstring"""







        SCREAMING_SNAKE_CASE :      Optional[int]       =   ()

        for resnet, attn in zip(self.resnets				,				self.attentions		):
              SCREAMING_SNAKE_CASE :      List[str]       =   resnet(UpperCAmelCase__				,				UpperCAmelCase__				,				deterministic=UpperCAmelCase__		)
              SCREAMING_SNAKE_CASE :      List[str]       =   attn(UpperCAmelCase__				,				UpperCAmelCase__				,				deterministic=UpperCAmelCase__		)
              output_states += (hidden_states,)

        if self.add_downsample:
              SCREAMING_SNAKE_CASE :      int       =   self.downsamplers_a(UpperCAmelCase__		)
              output_states += (hidden_states,)

        return hidden_states, output_states







class        a__       (   nn.Module       ):







  """simple docstring"""




  UpperCAmelCase__      : int
  UpperCAmelCase__      : int
  UpperCAmelCase__      : float			       =0.0
  UpperCAmelCase__      : int			       =1
  UpperCAmelCase__      : bool			       =True
  UpperCAmelCase__      : jnp.dtype			       =jnp.floataa


  def  _lowercase      (			self    : Dict		)					->Union[str, Any]:
        """simple docstring"""







        SCREAMING_SNAKE_CASE :      Dict       =   []

        for i in range(self.num_layers		):
              SCREAMING_SNAKE_CASE :      Union[str, Any]       =   self.in_channels if i == 0 else self.out_channels

              SCREAMING_SNAKE_CASE :      Union[str, Any]       =   FlaxResnetBlockaD(
                  in_channels=UpperCAmelCase__				,				out_channels=self.out_channels				,				dropout_prob=self.dropout				,				dtype=self.dtype				,				)
              resnets.append(UpperCAmelCase__		)
        SCREAMING_SNAKE_CASE :      Union[str, Any]       =   resnets

        if self.add_downsample:
              SCREAMING_SNAKE_CASE :      List[str]       =   FlaxDownsampleaD(self.out_channels				,				dtype=self.dtype		)




  def __call__(			self    : Tuple				,				UpperCAmelCase__    : int				,				UpperCAmelCase__    : Optional[Any]				,				UpperCAmelCase__    : Any=True		)					->Optional[Any]:
        """simple docstring"""







        SCREAMING_SNAKE_CASE :      Optional[int]       =   ()

        for resnet in self.resnets:
              SCREAMING_SNAKE_CASE :      List[str]       =   resnet(UpperCAmelCase__				,				UpperCAmelCase__				,				deterministic=UpperCAmelCase__		)
              output_states += (hidden_states,)

        if self.add_downsample:
              SCREAMING_SNAKE_CASE :      Dict       =   self.downsamplers_a(UpperCAmelCase__		)
              output_states += (hidden_states,)

        return hidden_states, output_states







class        a__       (   nn.Module       ):







  """simple docstring"""




  UpperCAmelCase__      : int
  UpperCAmelCase__      : int
  UpperCAmelCase__      : int
  UpperCAmelCase__      : float			       =0.0
  UpperCAmelCase__      : int			       =1
  UpperCAmelCase__      : int			       =1
  UpperCAmelCase__      : bool			       =True
  UpperCAmelCase__      : bool			       =False
  UpperCAmelCase__      : bool			       =False
  UpperCAmelCase__      : bool			       =False
  UpperCAmelCase__      : jnp.dtype			       =jnp.floataa


  def  _lowercase      (			self    : int		)					->Any:
        """simple docstring"""







        SCREAMING_SNAKE_CASE :      Dict       =   []
        SCREAMING_SNAKE_CASE :      List[Any]       =   []

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

              SCREAMING_SNAKE_CASE :      Tuple       =   FlaxResnetBlockaD(
                  in_channels=resnet_in_channels + res_skip_channels				,				out_channels=self.out_channels				,				dropout_prob=self.dropout				,				dtype=self.dtype				,				)
              resnets.append(UpperCAmelCase__		)

              SCREAMING_SNAKE_CASE :      Tuple       =   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(UpperCAmelCase__		)

        SCREAMING_SNAKE_CASE :      Optional[Any]       =   resnets
        SCREAMING_SNAKE_CASE :      List[Any]       =   attentions

        if self.add_upsample:
              SCREAMING_SNAKE_CASE :      Union[str, Any]       =   FlaxUpsampleaD(self.out_channels				,				dtype=self.dtype		)




  def __call__(			self    : Union[str, Any]				,				UpperCAmelCase__    : str				,				UpperCAmelCase__    : List[Any]				,				UpperCAmelCase__    : Any				,				UpperCAmelCase__    : Union[str, Any]				,				UpperCAmelCase__    : Dict=True		)					->Tuple:
        """simple docstring"""







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

              SCREAMING_SNAKE_CASE :      Tuple       =   resnet(UpperCAmelCase__				,				UpperCAmelCase__				,				deterministic=UpperCAmelCase__		)
              SCREAMING_SNAKE_CASE :      Tuple       =   attn(UpperCAmelCase__				,				UpperCAmelCase__				,				deterministic=UpperCAmelCase__		)

        if self.add_upsample:
              SCREAMING_SNAKE_CASE :      Tuple       =   self.upsamplers_a(UpperCAmelCase__		)

        return hidden_states







class        a__       (   nn.Module       ):







  """simple docstring"""




  UpperCAmelCase__      : int
  UpperCAmelCase__      : int
  UpperCAmelCase__      : int
  UpperCAmelCase__      : float			       =0.0
  UpperCAmelCase__      : int			       =1
  UpperCAmelCase__      : bool			       =True
  UpperCAmelCase__      : jnp.dtype			       =jnp.floataa


  def  _lowercase      (			self    : Union[str, Any]		)					->Optional[int]:
        """simple docstring"""







        SCREAMING_SNAKE_CASE :      Optional[Any]       =   []

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

              SCREAMING_SNAKE_CASE :      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(UpperCAmelCase__		)

        SCREAMING_SNAKE_CASE :      Dict       =   resnets

        if self.add_upsample:
              SCREAMING_SNAKE_CASE :      int       =   FlaxUpsampleaD(self.out_channels				,				dtype=self.dtype		)




  def __call__(			self    : Optional[Any]				,				UpperCAmelCase__    : List[str]				,				UpperCAmelCase__    : Dict				,				UpperCAmelCase__    : Tuple				,				UpperCAmelCase__    : Optional[int]=True		)					->Union[str, Any]:
        """simple docstring"""







        for resnet in self.resnets:
              # pop res hidden states
              SCREAMING_SNAKE_CASE :      Dict       =   res_hidden_states_tuple[-1]
              SCREAMING_SNAKE_CASE :      Any       =   res_hidden_states_tuple[:-1]
              SCREAMING_SNAKE_CASE :      Tuple       =   jnp.concatenate((hidden_states, res_hidden_states)				,				axis=-1		)

              SCREAMING_SNAKE_CASE :      Optional[int]       =   resnet(UpperCAmelCase__				,				UpperCAmelCase__				,				deterministic=UpperCAmelCase__		)

        if self.add_upsample:
              SCREAMING_SNAKE_CASE :      Optional[int]       =   self.upsamplers_a(UpperCAmelCase__		)

        return hidden_states







class        a__       (   nn.Module       ):







  """simple docstring"""




  UpperCAmelCase__      : int
  UpperCAmelCase__      : float			       =0.0
  UpperCAmelCase__      : int			       =1
  UpperCAmelCase__      : int			       =1
  UpperCAmelCase__      : bool			       =False
  UpperCAmelCase__      : bool			       =False
  UpperCAmelCase__      : jnp.dtype			       =jnp.floataa


  def  _lowercase      (			self    : Optional[Any]		)					->List[Any]:
        """simple docstring"""







        SCREAMING_SNAKE_CASE :      str       =   [
            FlaxResnetBlockaD(
                in_channels=self.in_channels				,				out_channels=self.in_channels				,				dropout_prob=self.dropout				,				dtype=self.dtype				,				)
        ]

        SCREAMING_SNAKE_CASE :      List[str]       =   []

        for _ in range(self.num_layers		):
              SCREAMING_SNAKE_CASE :      Tuple       =   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(UpperCAmelCase__		)

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

        SCREAMING_SNAKE_CASE :      int       =   resnets
        SCREAMING_SNAKE_CASE :      Tuple       =   attentions




  def __call__(			self    : Union[str, Any]				,				UpperCAmelCase__    : List[str]				,				UpperCAmelCase__    : Union[str, Any]				,				UpperCAmelCase__    : Any				,				UpperCAmelCase__    : str=True		)					->Optional[Any]:
        """simple docstring"""







        SCREAMING_SNAKE_CASE :      Tuple       =   self.resnets[0](UpperCAmelCase__				,				UpperCAmelCase__		)
        for attn, resnet in zip(self.attentions				,				self.resnets[1:]		):
              SCREAMING_SNAKE_CASE :      Optional[int]       =   attn(UpperCAmelCase__				,				UpperCAmelCase__				,				deterministic=UpperCAmelCase__		)
              SCREAMING_SNAKE_CASE :      List[str]       =   resnet(UpperCAmelCase__				,				UpperCAmelCase__				,				deterministic=UpperCAmelCase__		)

        return hidden_states




