

"""simple docstring"""




from dataclasses import dataclass
from typing import Dict, Optional, Union

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

from ..configuration_utils import ConfigMixin, register_to_config
from ..utils import BaseOutput
from .attention import BasicTransformerBlock
from .attention_processor import AttentionProcessor, AttnProcessor
from .embeddings import TimestepEmbedding, Timesteps
from .modeling_utils import ModelMixin







@dataclass
class   __lowercase      (     lowercase_       ):
  snake_case_												=							42







class   __lowercase      (     lowercase_						,      lowercase_       ):

  @register_to_config
  def __init__(	self		:      Any       ,A		:      Union[str, Any] = 32       ,A		:      List[str] = 64       ,A		:      str = 20       ,A		:      str = 768       ,A		:      Dict=77       ,A		:      str=4       ,A		:      Optional[Any] = 0.0       ,A		:      List[str] = "silu"       ,A		:      List[Any] = None       ,A		:      Any = None       ,A		:      str = "linear"       ,A		:      Any = "prd"       ,A		:      Any = None       ,A		:      List[str] = None       ,A		:      Tuple = None       ,):






   '''simple docstring'''
   super().__init__()
   UpperCAmelCase__       :	List[str]        =       num_attention_heads
   UpperCAmelCase__       :	List[Any]        =       attention_head_dim
   UpperCAmelCase__       :	Optional[int]        =       num_attention_heads * attention_head_dim
   UpperCAmelCase__       :	Optional[int]        =       additional_embeddings

   UpperCAmelCase__       :	Optional[Any]        =       time_embed_dim or inner_dim
   UpperCAmelCase__       :	Tuple        =       embedding_proj_dim or embedding_dim
   UpperCAmelCase__       :	List[str]        =       clip_embed_dim or embedding_dim

   UpperCAmelCase__       :	Optional[int]        =       Timesteps(a__       ,a__       ,0     )
   UpperCAmelCase__       :	Dict        =       TimestepEmbedding(a__       ,a__       ,out_dim=a__       ,act_fn=a__     )

   UpperCAmelCase__       :	Union[str, Any]        =       nn.Linear(a__       ,a__     )

   if embedding_proj_norm_type is None:
    UpperCAmelCase__       :	int        =       None
   elif embedding_proj_norm_type == "layer":
    UpperCAmelCase__       :	List[Any]        =       nn.LayerNorm(a__     )
   else:
    raise ValueError(f"unsupported embedding_proj_norm_type: {embedding_proj_norm_type}"     )

   UpperCAmelCase__       :	Tuple        =       nn.Linear(a__       ,a__     )

   if encoder_hid_proj_type is None:
    UpperCAmelCase__       :	List[Any]        =       None
   elif encoder_hid_proj_type == "linear":
    UpperCAmelCase__       :	int        =       nn.Linear(a__       ,a__     )
   else:
    raise ValueError(f"unsupported encoder_hid_proj_type: {encoder_hid_proj_type}"     )

   UpperCAmelCase__       :	int        =       nn.Parameter(torch.zeros(1       ,num_embeddings + additional_embeddings       ,a__     )     )

   if added_emb_type == "prd":
    UpperCAmelCase__       :	Tuple        =       nn.Parameter(torch.zeros(1       ,1       ,a__     )     )
   elif added_emb_type is None:
    UpperCAmelCase__       :	Tuple        =       None
   else:
    raise ValueError(
        f"`added_emb_type`: {added_emb_type} is not supported. Make sure to choose one of `\'prd\'` or `None`."     )

   UpperCAmelCase__       :	Tuple        =       nn.ModuleList(
       [
           BasicTransformerBlock(
               a__       ,a__       ,a__       ,dropout=a__       ,activation_fn="""gelu"""       ,attention_bias=a__       ,)
           for d in range(a__     )
       ]     )

   if norm_in_type == "layer":
    UpperCAmelCase__       :	Any        =       nn.LayerNorm(a__     )
   elif norm_in_type is None:
    UpperCAmelCase__       :	str        =       None
   else:
    raise ValueError(f"Unsupported norm_in_type: {norm_in_type}."     )

   UpperCAmelCase__       :	int        =       nn.LayerNorm(a__     )

   UpperCAmelCase__       :	Optional[Any]        =       nn.Linear(a__       ,a__     )

   UpperCAmelCase__       :	Any        =       torch.full(
       [num_embeddings + additional_embeddings, num_embeddings + additional_embeddings]       ,-10_000.0     )
   causal_attention_mask.triu_(1     )
   UpperCAmelCase__       :	Dict        =       causal_attention_mask[None, ...]
   self.register_buffer("""causal_attention_mask"""       ,a__       ,persistent=a__     )

   UpperCAmelCase__       :	List[str]        =       nn.Parameter(torch.zeros(1       ,a__     )     )
   UpperCAmelCase__       :	Optional[int]        =       nn.Parameter(torch.zeros(1       ,a__     )     )

  @property
  # Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.attn_processors
  def  __lowercase							(	self		:      int     ):






   '''simple docstring'''
   UpperCAmelCase__       :	List[Any]        =       {}

   def fn_recursive_add_processors(A		:      Dict       ,A		:      Dict       ,A		:      Any     ):
    if hasattr(a__       ,"""set_processor"""     ):
     UpperCAmelCase__       :	int        =       module.processor

    for sub_name, child in module.named_children():
     fn_recursive_add_processors(f"{name}.{sub_name}"       ,a__       ,a__     )

    return processors

   for name, module in self.named_children():
    fn_recursive_add_processors(a__       ,a__       ,a__     )

   return processors

  def  __lowercase							(	self		:      str       ,A		:      Any     ):






   '''simple docstring'''
   UpperCAmelCase__       :	str        =       len(self.attn_processors.keys()     )

   if isinstance(a__       ,a__     ) and len(a__     ) != count:
    raise ValueError(
        f"A dict of processors was passed, but the number of processors {len(a__     )} does not match the"
        f" number of attention layers: {count}. Please make sure to pass {count} processor classes."     )

   def fn_recursive_attn_processor(A		:      Optional[Any]       ,A		:      Dict       ,A		:      List[Any]     ):
    if hasattr(a__       ,"""set_processor"""     ):
     if not isinstance(a__       ,a__     ):
      module.set_processor(a__     )
     else:
      module.set_processor(processor.pop(f"{name}.processor"     )     )

    for sub_name, child in module.named_children():
     fn_recursive_attn_processor(f"{name}.{sub_name}"       ,a__       ,a__     )

   for name, module in self.named_children():
    fn_recursive_attn_processor(a__       ,a__       ,a__     )

  def  __lowercase							(	self		:      List[Any]     ):






   '''simple docstring'''
   self.set_attn_processor(AttnProcessor()     )

  def  __lowercase							(	self		:      Tuple       ,A		:      Union[str, Any]       ,A		:      int       ,A		:      List[str]       ,A		:      Any = None       ,A		:      List[Any] = None       ,A		:      Optional[int] = True       ,):






   '''simple docstring'''
   UpperCAmelCase__       :	List[Any]        =       hidden_states.shape[0]

   UpperCAmelCase__       :	Union[str, Any]        =       timestep
   if not torch.is_tensor(a__     ):
    UpperCAmelCase__       :	List[Any]        =       torch.tensor([timesteps]       ,dtype=torch.long       ,device=hidden_states.device     )
   elif torch.is_tensor(a__     ) and len(timesteps.shape     ) == 0:
    UpperCAmelCase__       :	Any        =       timesteps[None].to(hidden_states.device     )

   # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
   UpperCAmelCase__       :	str        =       timesteps * torch.ones(a__       ,dtype=timesteps.dtype       ,device=timesteps.device     )

   UpperCAmelCase__       :	Optional[int]        =       self.time_proj(a__     )

   # timesteps does not contain any weights and will always return f32 tensors
   # but time_embedding might be fp16, so we need to cast here.
   UpperCAmelCase__       :	List[str]        =       timesteps_projected.to(dtype=self.dtype     )
   UpperCAmelCase__       :	Tuple        =       self.time_embedding(a__     )

   if self.embedding_proj_norm is not None:
    UpperCAmelCase__       :	Dict        =       self.embedding_proj_norm(a__     )

   UpperCAmelCase__       :	Optional[Any]        =       self.embedding_proj(a__     )
   if self.encoder_hidden_states_proj is not None and encoder_hidden_states is not None:
    UpperCAmelCase__       :	Any        =       self.encoder_hidden_states_proj(a__     )
   elif self.encoder_hidden_states_proj is not None and encoder_hidden_states is None:
    raise ValueError("""`encoder_hidden_states_proj` requires `encoder_hidden_states` to be set"""     )

   UpperCAmelCase__       :	List[str]        =       self.proj_in(a__     )

   UpperCAmelCase__       :	Any        =       self.positional_embedding.to(hidden_states.dtype     )

   UpperCAmelCase__       :	Any        =       []
   UpperCAmelCase__       :	List[Any]        =       0

   if encoder_hidden_states is not None:
    additional_embeds.append(a__     )
    additional_embeddings_len += encoder_hidden_states.shape[1]

   if len(proj_embeddings.shape     ) == 2:
    UpperCAmelCase__       :	str        =       proj_embeddings[:, None, :]

   if len(hidden_states.shape     ) == 2:
    UpperCAmelCase__       :	Optional[Any]        =       hidden_states[:, None, :]

   UpperCAmelCase__       :	str        =       additional_embeds + [
       proj_embeddings,
       time_embeddings[:, None, :],
       hidden_states,
   ]

   if self.prd_embedding is not None:
    UpperCAmelCase__       :	List[str]        =       self.prd_embedding.to(hidden_states.dtype     ).expand(a__       ,-1       ,-1     )
    additional_embeds.append(a__     )

   UpperCAmelCase__       :	int        =       torch.cat(
       a__       ,dim=1       ,)

   # Allow positional_embedding to not include the `addtional_embeddings` and instead pad it with zeros for these additional tokens
   UpperCAmelCase__       :	Dict        =       additional_embeddings_len + proj_embeddings.shape[1] + 1
   if positional_embeddings.shape[1] < hidden_states.shape[1]:
    UpperCAmelCase__       :	int        =       F.pad(
        a__       ,(
            0,
            0,
            additional_embeddings_len,
            self.prd_embedding.shape[1] if self.prd_embedding is not None else 0,
        )       ,value=0.0       ,)

   UpperCAmelCase__       :	Any        =       hidden_states + positional_embeddings

   if attention_mask is not None:
    UpperCAmelCase__       :	Tuple        =       (1 - attention_mask.to(hidden_states.dtype     )) * -10_000.0
    UpperCAmelCase__       :	List[str]        =       F.pad(a__       ,(0, self.additional_embeddings)       ,value=0.0     )
    UpperCAmelCase__       :	str        =       (attention_mask[:, None, :] + self.causal_attention_mask).to(hidden_states.dtype     )
    UpperCAmelCase__       :	Tuple        =       attention_mask.repeat_interleave(self.config.num_attention_heads       ,dim=0     )

   if self.norm_in is not None:
    UpperCAmelCase__       :	int        =       self.norm_in(a__     )

   for block in self.transformer_blocks:
    UpperCAmelCase__       :	Union[str, Any]        =       block(a__       ,attention_mask=a__     )

   UpperCAmelCase__       :	Optional[Any]        =       self.norm_out(a__     )

   if self.prd_embedding is not None:
    UpperCAmelCase__       :	str        =       hidden_states[:, -1]
   else:
    UpperCAmelCase__       :	Dict        =       hidden_states[:, additional_embeddings_len:]

   UpperCAmelCase__       :	Tuple        =       self.proj_to_clip_embeddings(a__     )

   if not return_dict:
    return (predicted_image_embedding,)

   return PriorTransformerOutput(predicted_image_embedding=a__     )




  def  __lowercase							(	self		:      Dict       ,A		:      Union[str, Any]     ):






   '''simple docstring'''
   UpperCAmelCase__       :	str        =       (prior_latents * self.clip_std) + self.clip_mean
   return prior_latents













import json
import os
import tempfile

import transformers

import datasets
from utils import generate_example_dataset, get_duration


lowerCamelCase__       		=  50_0000

lowerCamelCase__		,      lowerCamelCase__       		=  os.path.split(__file__)
lowerCamelCase__       		=  os.path.join(RESULTS_BASEPATH, """results""", RESULTS_FILENAME.replace(""".py""", """.json"""))






@get_duration
def 	UpperCamelCase	(							snake_case__  : datasets.Dataset						,**snake_case__  : Any	):




     '''simple docstring'''






     __snake_case      :List[str] 							=    dataset.map(**snake_case__	)






@get_duration
def 	UpperCamelCase	(							snake_case__  : datasets.Dataset						,**snake_case__  : Dict	):




     '''simple docstring'''






     __snake_case      :int 							=    dataset.filter(**snake_case__	)






def 	UpperCamelCase	(							):




     '''simple docstring'''






     __snake_case      :str 							=    {"""num examples""": SPEED_TEST_N_EXAMPLES}
     with tempfile.TemporaryDirectory() as tmp_dir:
          __snake_case      :Dict 							=    datasets.Features({"""text""": datasets.Value("""string"""	), """numbers""": datasets.Value("""float32"""	)}	)
          __snake_case      :Optional[int] 							=    generate_example_dataset(
              os.path.join(snake_case__						,"""dataset.arrow"""	)						,snake_case__						,num_examples=snake_case__	)

          __snake_case      :str 							=    transformers.AutoTokenizer.from_pretrained("""bert-base-cased"""						,use_fast=snake_case__	)

          def tokenize(snake_case__  : Optional[int]	):
               return tokenizer(examples["""text"""]	)

          __snake_case      :Optional[Any] 							=    map(snake_case__	)

          __snake_case      :Optional[Any] 							=    map(snake_case__						,batched=snake_case__	)

          __snake_case      :Optional[Any] 							=    map(snake_case__						,function=lambda snake_case__	: None						,batched=snake_case__	)

          with dataset.formatted_as(type="""numpy"""	):
               __snake_case      :int 							=    map(snake_case__						,function=lambda snake_case__	: None						,batched=snake_case__	)

          with dataset.formatted_as(type="""pandas"""	):
               __snake_case      :List[Any] 							=    map(snake_case__						,function=lambda snake_case__	: None						,batched=snake_case__	)

          with dataset.formatted_as(type="""torch"""						,columns="""numbers"""	):
               __snake_case      :str 							=    map(snake_case__						,function=lambda snake_case__	: None						,batched=snake_case__	)

          with dataset.formatted_as(type="""tensorflow"""						,columns="""numbers"""	):
               __snake_case      :Optional[int] 							=    map(snake_case__						,function=lambda snake_case__	: None						,batched=snake_case__	)

          __snake_case      :Dict 							=    map(snake_case__						,function=snake_case__						,batched=snake_case__	)

          __snake_case      :Tuple 							=    filter(snake_case__	)

          # Activate later when tokenizer support batched inputs
          # with dataset.formatted_as(type='numpy'):
          #     times[func.__name__ + " fast-tokenizer batched numpy"] = func(dataset, function=tokenize, batched=True)

     with open(snake_case__						,"""wb"""	) as f:
          f.write(json.dumps(snake_case__	).encode("""utf-8"""	)	)


if __name__ == "__main__":  # useful to run the profiler
     benchmark_map_filter()


