


'''simple docstring'''







from typing import Optional, Tuple, Union

import flax
import flax.linen as nn
import jax
import jax.numpy as jnp
from flax.core.frozen_dict import FrozenDict

from ..configuration_utils import ConfigMixin, flax_register_to_config
from ..utils import BaseOutput
from .embeddings_flax import FlaxTimestepEmbedding, FlaxTimesteps
from .modeling_flax_utils import FlaxModelMixin
from .unet_ad_blocks_flax import (
    FlaxCrossAttnDownBlockaD,
    FlaxDownBlockaD,
    FlaxUNetMidBlockaDCrossAttn,
)


@flax.struct.dataclass
class  lowerCAmelCase_						(       __a      ):





  """simple docstring"""



  a_       :jnp.ndarray
  a_       :jnp.ndarray


class  lowerCAmelCase_						(       nn.Module      ):





  """simple docstring"""



  a_       :int
  a_       :Tuple[int]												=(16, 32, 96, 256)
  a_       :jnp.dtype												=jnp.floataa

  def __a						(					self     :							Dict						):





       '''simple docstring'''

       __a       =   nn.Conv(
           self.block_out_channels[0]				,							kernel_size=(3, 3)				,							padding=((1, 1), (1, 1))				,							dtype=self.dtype				,							)

       __a       =   []
       for i in range(len(self.block_out_channels						) - 1						):
            __a       =   self.block_out_channels[i]
            __a       =   self.block_out_channels[i + 1]
            __a       =   nn.Conv(
                lowerCAmelCase_				,							kernel_size=(3, 3)				,							padding=((1, 1), (1, 1))				,							dtype=self.dtype				,							)
            blocks.append(lowerCAmelCase_						)
            __a       =   nn.Conv(
                lowerCAmelCase_				,							kernel_size=(3, 3)				,							strides=(2, 2)				,							padding=((1, 1), (1, 1))				,							dtype=self.dtype				,							)
            blocks.append(lowerCAmelCase_						)
       __a       =   blocks

       __a       =   nn.Conv(
           self.conditioning_embedding_channels				,							kernel_size=(3, 3)				,							padding=((1, 1), (1, 1))				,							kernel_init=nn.initializers.zeros_init()				,							bias_init=nn.initializers.zeros_init()				,							dtype=self.dtype				,							)


  def __call__(					self     :							List[Any]				,							SCREAMING_SNAKE_CASE__     :							str						):





       '''simple docstring'''

       __a       =   self.conv_in(lowerCAmelCase_						)
       __a       =   nn.silu(lowerCAmelCase_						)

       for block in self.blocks:
            __a       =   block(lowerCAmelCase_						)
            __a       =   nn.silu(lowerCAmelCase_						)

       __a       =   self.conv_out(lowerCAmelCase_						)

       return embedding


@flax_register_to_config
class  lowerCAmelCase_						(       nn.Module		,    __a		,    __a      ):





  """simple docstring"""



  a_       :int												=32
  a_       :int												=4
  a_       :Tuple[str]												=(
      "CrossAttnDownBlock2D",
      "CrossAttnDownBlock2D",
      "CrossAttnDownBlock2D",
      "DownBlock2D",
  )
  a_       :Union[bool, Tuple[bool]]												=False
  a_       :Tuple[int]												=(320, 640, 1280, 1280)
  a_       :int												=2
  a_       :Union[int, Tuple[int]]												=8
  a_       :Optional[Union[int, Tuple[int]]]												=None
  a_       :int												=1280
  a_       :float												=0.0
  a_       :bool												=False
  a_       :jnp.dtype												=jnp.floataa
  a_       :bool												=True
  a_       :int												=0
  a_       :str												="rgb"
  a_       :Tuple[int]												=(16, 32, 96, 256)

  def __a						(					self     :							Dict				,							SCREAMING_SNAKE_CASE__     :							Any						):





       '''simple docstring'''

       __a       =   (1, self.in_channels, self.sample_size, self.sample_size)
       __a       =   jnp.zeros(lowerCAmelCase_				,							dtype=jnp.floataa						)
       __a       =   jnp.ones((1,)				,							dtype=jnp.intaa						)
       __a       =   jnp.zeros((1, 1, self.cross_attention_dim)				,							dtype=jnp.floataa						)
       __a       =   (1, 3, self.sample_size * 8, self.sample_size * 8)
       __a       =   jnp.zeros(lowerCAmelCase_				,							dtype=jnp.floataa						)

       __a ,						__a       =   jax.random.split(lowerCAmelCase_						)
       __a       =   {"""params""": params_rng, """dropout""": dropout_rng}

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

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





       '''simple docstring'''

       __a       =   self.block_out_channels
       __a       =   block_out_channels[0] * 4

       # If `num_attention_heads` is not defined (which is the case for most models)
       # it will default to `attention_head_dim`. This looks weird upon first reading it and it is.
       # The reason for this behavior is to correct for incorrectly named variables that were introduced
       # when this library was created. The incorrect naming was only discovered much later in https://github.com/huggingface/diffusers/issues/2011#issuecomment-1547958131
       # Changing `attention_head_dim` to `num_attention_heads` for 40,000+ configurations is too backwards breaking
       # which is why we correct for the naming here.
       __a       =   self.num_attention_heads or self.attention_head_dim

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

       # time
       __a       =   FlaxTimesteps(
           block_out_channels[0]				,							flip_sin_to_cos=self.flip_sin_to_cos				,							freq_shift=self.config.freq_shift						)
       __a       =   FlaxTimestepEmbedding(lowerCAmelCase_				,							dtype=self.dtype						)

       __a       =   FlaxControlNetConditioningEmbedding(
           conditioning_embedding_channels=block_out_channels[0]				,							block_out_channels=self.conditioning_embedding_out_channels				,							)

       __a       =   self.only_cross_attention
       if isinstance(lowerCAmelCase_				,							lowerCAmelCase_						):
            __a       =   (only_cross_attention,) * len(self.down_block_types						)

       if isinstance(lowerCAmelCase_				,							lowerCAmelCase_						):
            __a       =   (num_attention_heads,) * len(self.down_block_types						)

       # down
       __a       =   []
       __a       =   []

       __a       =   block_out_channels[0]

       __a       =   nn.Conv(
           lowerCAmelCase_				,							kernel_size=(1, 1)				,							padding="""VALID"""				,							kernel_init=nn.initializers.zeros_init()				,							bias_init=nn.initializers.zeros_init()				,							dtype=self.dtype				,							)
       controlnet_down_blocks.append(lowerCAmelCase_						)

       for i, down_block_type in enumerate(self.down_block_types						):
            __a       =   output_channel
            __a       =   block_out_channels[i]
            __a       =   i == len(lowerCAmelCase_						) - 1

            if down_block_type == "CrossAttnDownBlock2D":
                 __a       =   FlaxCrossAttnDownBlockaD(
                     in_channels=lowerCAmelCase_				,							out_channels=lowerCAmelCase_				,							dropout=self.dropout				,							num_layers=self.layers_per_block				,							num_attention_heads=num_attention_heads[i]				,							add_downsample=not is_final_block				,							use_linear_projection=self.use_linear_projection				,							only_cross_attention=only_cross_attention[i]				,							dtype=self.dtype				,							)
            else:
                 __a       =   FlaxDownBlockaD(
                     in_channels=lowerCAmelCase_				,							out_channels=lowerCAmelCase_				,							dropout=self.dropout				,							num_layers=self.layers_per_block				,							add_downsample=not is_final_block				,							dtype=self.dtype				,							)

            down_blocks.append(lowerCAmelCase_						)

            for _ in range(self.layers_per_block						):
                 __a       =   nn.Conv(
                     lowerCAmelCase_				,							kernel_size=(1, 1)				,							padding="""VALID"""				,							kernel_init=nn.initializers.zeros_init()				,							bias_init=nn.initializers.zeros_init()				,							dtype=self.dtype				,							)
                 controlnet_down_blocks.append(lowerCAmelCase_						)

            if not is_final_block:
                 __a       =   nn.Conv(
                     lowerCAmelCase_				,							kernel_size=(1, 1)				,							padding="""VALID"""				,							kernel_init=nn.initializers.zeros_init()				,							bias_init=nn.initializers.zeros_init()				,							dtype=self.dtype				,							)
                 controlnet_down_blocks.append(lowerCAmelCase_						)

       __a       =   down_blocks
       __a       =   controlnet_down_blocks

       # mid
       __a       =   block_out_channels[-1]
       __a       =   FlaxUNetMidBlockaDCrossAttn(
           in_channels=lowerCAmelCase_				,							dropout=self.dropout				,							num_attention_heads=num_attention_heads[-1]				,							use_linear_projection=self.use_linear_projection				,							dtype=self.dtype				,							)

       __a       =   nn.Conv(
           lowerCAmelCase_				,							kernel_size=(1, 1)				,							padding="""VALID"""				,							kernel_init=nn.initializers.zeros_init()				,							bias_init=nn.initializers.zeros_init()				,							dtype=self.dtype				,							)


  def __call__(					self     :							str				,							SCREAMING_SNAKE_CASE__     :							Union[str, Any]				,							SCREAMING_SNAKE_CASE__     :							Dict				,							SCREAMING_SNAKE_CASE__     :							Tuple				,							SCREAMING_SNAKE_CASE__     :							List[Any]				,							SCREAMING_SNAKE_CASE__     :							Optional[Any] = 1.0				,							SCREAMING_SNAKE_CASE__     :							Any = True				,							SCREAMING_SNAKE_CASE__     :							Optional[Any] = False				,							):





       '''simple docstring'''

       __a       =   self.controlnet_conditioning_channel_order
       if channel_order == "bgr":
            __a       =   jnp.flip(lowerCAmelCase_				,							axis=1						)

       # 1. time
       if not isinstance(lowerCAmelCase_				,							jnp.ndarray						):
            __a       =   jnp.array([timesteps]				,							dtype=jnp.intaa						)
       elif isinstance(lowerCAmelCase_				,							jnp.ndarray						) and len(timesteps.shape						) == 0:
            __a       =   timesteps.astype(dtype=jnp.floataa						)
            __a       =   jnp.expand_dims(lowerCAmelCase_				,							0						)

       __a       =   self.time_proj(lowerCAmelCase_						)
       __a       =   self.time_embedding(lowerCAmelCase_						)

       # 2. pre-process
       __a       =   jnp.transpose(lowerCAmelCase_				,							(0, 2, 3, 1)						)
       __a       =   self.conv_in(lowerCAmelCase_						)

       __a       =   jnp.transpose(lowerCAmelCase_				,							(0, 2, 3, 1)						)
       __a       =   self.controlnet_cond_embedding(lowerCAmelCase_						)
       sample += controlnet_cond

       # 3. down
       __a       =   (sample,)
       for down_block in self.down_blocks:
            if isinstance(lowerCAmelCase_				,							lowerCAmelCase_						):
                 __a ,						__a       =   down_block(lowerCAmelCase_				,							lowerCAmelCase_				,							lowerCAmelCase_				,							deterministic=not train						)
            else:
                 __a ,						__a       =   down_block(lowerCAmelCase_				,							lowerCAmelCase_				,							deterministic=not train						)
            down_block_res_samples += res_samples

       # 4. mid
       __a       =   self.mid_block(lowerCAmelCase_				,							lowerCAmelCase_				,							lowerCAmelCase_				,							deterministic=not train						)

       # 5. contronet blocks
       __a       =   ()
       for down_block_res_sample, controlnet_block in zip(lowerCAmelCase_				,							self.controlnet_down_blocks						):
            __a       =   controlnet_block(lowerCAmelCase_						)
            controlnet_down_block_res_samples += (down_block_res_sample,)

       __a       =   controlnet_down_block_res_samples

       __a       =   self.controlnet_mid_block(lowerCAmelCase_						)

       # 6. scaling
       __a       =   [sample * conditioning_scale for sample in down_block_res_samples]
       mid_block_res_sample *= conditioning_scale

       if not return_dict:
            return (down_block_res_samples, mid_block_res_sample)

       return FlaxControlNetOutput(
           down_block_res_samples=lowerCAmelCase_				,							mid_block_res_sample=lowerCAmelCase_						)










'''simple docstring'''







import argparse
from typing import Dict

import tensorflow as tf
import torch
from tqdm import tqdm

from transformers import BigBirdPegasusConfig, BigBirdPegasusForConditionalGeneration


SCREAMING_SNAKE_CASE_      =			[
    # tf -> hf
    ('/', '.'),
    ('layer_', 'layers.'),
    ('kernel', 'weight'),
    ('beta', 'bias'),
    ('gamma', 'weight'),
    ('pegasus', 'model'),
]
SCREAMING_SNAKE_CASE_      =			[
    ('.output.dense', '.fc2'),
    ('intermediate.LayerNorm', 'final_layer_norm'),
    ('intermediate.dense', 'fc1'),
]

SCREAMING_SNAKE_CASE_      =			(
    INIT_COMMON
    + [
        ('attention.self.LayerNorm', 'self_attn_layer_norm'),
        ('attention.output.dense', 'self_attn.out_proj'),
        ('attention.self', 'self_attn'),
        ('attention.encdec.LayerNorm', 'encoder_attn_layer_norm'),
        ('attention.encdec_output.dense', 'encoder_attn.out_proj'),
        ('attention.encdec', 'encoder_attn'),
        ('key', 'k_proj'),
        ('value', 'v_proj'),
        ('query', 'q_proj'),
        ('decoder.LayerNorm', 'decoder.layernorm_embedding'),
    ]
    + END_COMMON
)

SCREAMING_SNAKE_CASE_      =			(
    INIT_COMMON
    + [
        ('embeddings.word_embeddings', 'shared.weight'),
        ('embeddings.position_embeddings', 'embed_positions.weight'),
        ('attention.self.LayerNorm', 'self_attn_layer_norm'),
        ('attention.output.dense', 'self_attn.output'),
        ('attention.self', 'self_attn.self'),
        ('encoder.LayerNorm', 'encoder.layernorm_embedding'),
    ]
    + END_COMMON
)

SCREAMING_SNAKE_CASE_      =			[
    'encdec/key/bias',
    'encdec/query/bias',
    'encdec/value/bias',
    'self/key/bias',
    'self/query/bias',
    'self/value/bias',
    'encdec_output/dense/bias',
    'attention/output/dense/bias',
]




def        __lowercase						(      __SCREAMING_SNAKE_CASE				,  __SCREAMING_SNAKE_CASE		)				->					int:
     """simple docstring"""




     for tf_name, hf_name in patterns:
          __a       =   k.replace(__SCREAMING_SNAKE_CASE				,  __SCREAMING_SNAKE_CASE		)
     return k




def        __lowercase						(      __SCREAMING_SNAKE_CASE				,  __SCREAMING_SNAKE_CASE		)				->					BigBirdPegasusForConditionalGeneration:
     """simple docstring"""




     __a       =   BigBirdPegasusConfig(**__SCREAMING_SNAKE_CASE		)
     __a       =   BigBirdPegasusForConditionalGeneration(__SCREAMING_SNAKE_CASE		)
     __a       =   torch_model.state_dict()
     __a       =   {}

     # separating decoder weights
     __a       =   {k: tf_weights[k] for k in tf_weights if k.startswith("""pegasus/decoder"""		)}
     __a       =   {k: tf_weights[k] for k in tf_weights if not k.startswith("""pegasus/decoder"""		)}

     for k, v in tqdm(decoder_weights.items()				,  """tf -> hf conversion"""		):
          __a       =   [k.endswith(__SCREAMING_SNAKE_CASE		) for ending in KEYS_TO_IGNORE]
          if any(__SCREAMING_SNAKE_CASE		):
               continue
          __a       =   DECODER_PATTERNS
          __a       =   rename_state_dict_key(__SCREAMING_SNAKE_CASE				,  __SCREAMING_SNAKE_CASE		)
          if new_k not in state_dict:
               raise ValueError(F'''could not find new key {new_k} in state dict. (converted from {k})'''		)
          if any(True if i in k else False for i in ["""dense""", """query""", """key""", """value"""]		):
               __a       =   v.T
          __a       =   torch.from_numpy(__SCREAMING_SNAKE_CASE		)
          assert v.shape == state_dict[new_k].shape, F'''{new_k}, {k}, {v.shape}, {state_dict[new_k].shape}'''

     for k, v in tqdm(remaining_weights.items()				,  """tf -> hf conversion"""		):
          __a       =   [k.endswith(__SCREAMING_SNAKE_CASE		) for ending in KEYS_TO_IGNORE]
          if any(__SCREAMING_SNAKE_CASE		):
               continue
          __a       =   REMAINING_PATTERNS
          __a       =   rename_state_dict_key(__SCREAMING_SNAKE_CASE				,  __SCREAMING_SNAKE_CASE		)
          if new_k not in state_dict and k != "pegasus/embeddings/position_embeddings":
               raise ValueError(F'''could not find new key {new_k} in state dict. (converted from {k})'''		)
          if any(True if i in k else False for i in ["""dense""", """query""", """key""", """value"""]		):
               __a       =   v.T
          __a       =   torch.from_numpy(__SCREAMING_SNAKE_CASE		)
          if k != "pegasus/embeddings/position_embeddings":
               assert v.shape == state_dict[new_k].shape, F'''{new_k}, {k}, {v.shape}, {state_dict[new_k].shape}'''

     __a       =   mapping["""model.embed_positions.weight"""]
     __a       =   mapping.pop("""model.embed_positions.weight"""		)
     __a ,						__a       =   torch_model.load_state_dict(__SCREAMING_SNAKE_CASE				,  strict=__SCREAMING_SNAKE_CASE		)
     __a       =   [
         k
         for k in missing
         if k
         not in [
             """final_logits_bias""",
             """model.encoder.embed_tokens.weight""",
             """model.decoder.embed_tokens.weight""",
             """lm_head.weight""",
         ]
     ]
     assert unexpected_missing == [], F'''no matches found for the following torch keys {unexpected_missing}'''
     assert extra == [], F'''no matches found for the following tf keys {extra}'''
     return torch_model




def        __lowercase						(      __SCREAMING_SNAKE_CASE		)				->					Dict:
     """simple docstring"""




     __a       =   tf.train.list_variables(__SCREAMING_SNAKE_CASE		)
     __a       =   {}
     __a       =   ["""global_step"""]
     for name, shape in tqdm(__SCREAMING_SNAKE_CASE				,  desc="""converting tf checkpoint to dict"""		):
          __a       =   any(pat in name for pat in ignore_name		)
          if skip_key:
               continue
          __a       =   tf.train.load_variable(__SCREAMING_SNAKE_CASE				,  __SCREAMING_SNAKE_CASE		)
          __a       =   array
     return tf_weights




def        __lowercase						(      __SCREAMING_SNAKE_CASE				,  __SCREAMING_SNAKE_CASE				,  __SCREAMING_SNAKE_CASE		)				->					Tuple:
     """simple docstring"""




     __a       =   get_tf_weights_as_numpy(__SCREAMING_SNAKE_CASE		)
     __a       =   convert_bigbird_pegasus(__SCREAMING_SNAKE_CASE				,  __SCREAMING_SNAKE_CASE		)
     torch_model.save_pretrained(__SCREAMING_SNAKE_CASE		)


if __name__ == "__main__":
 SCREAMING_SNAKE_CASE_      =			argparse.ArgumentParser()
 parser.add_argument('--tf_ckpt_path', type=str, help='passed to tf.train.list_variables')
 parser.add_argument('--save_dir', default=None, type=str, help='Path to the output PyTorch model.')
 SCREAMING_SNAKE_CASE_      =			parser.parse_args()
 SCREAMING_SNAKE_CASE_      =			{}
 convert_bigbird_pegasus_ckpt_to_pytorch(args.tf_ckpt_path, args.save_dir, config_update=config_update)






