




from typing import Optional

from torch import nn

from .transformer_ad import TransformeraDModel, TransformeraDModelOutput




class       __lowerCAmelCase				(							nn.Module):






       def __init__(      self:    Union[str, Any]				,    _lowerCAmelCase:    int = 16				,    _lowerCAmelCase:    int = 88				,    _lowerCAmelCase:    Optional[int] = None				,    _lowerCAmelCase:    int = 1				,    _lowerCAmelCase:    float = 0.0				,    _lowerCAmelCase:    int = 32				,    _lowerCAmelCase:    Optional[int] = None				,    _lowerCAmelCase:    bool = False				,    _lowerCAmelCase:    Optional[int] = None				,    _lowerCAmelCase:    Optional[int] = None				,    _lowerCAmelCase:    str = "geglu"				,    _lowerCAmelCase:    Optional[int] = None				,    ):
         super().__init__()
         lowercase			:Optional[int]							=  nn.ModuleList(
             [
                 TransformeraDModel(
                     num_attention_heads=_lowerCAmelCase				,    attention_head_dim=_lowerCAmelCase				,    in_channels=_lowerCAmelCase				,    num_layers=_lowerCAmelCase				,    dropout=_lowerCAmelCase				,    norm_num_groups=_lowerCAmelCase				,    cross_attention_dim=_lowerCAmelCase				,    attention_bias=_lowerCAmelCase				,    sample_size=_lowerCAmelCase				,    num_vector_embeds=_lowerCAmelCase				,    activation_fn=_lowerCAmelCase				,    num_embeds_ada_norm=_lowerCAmelCase				,    )
                 for _ in range(2 )
             ] )

         # Variables that can be set by a pipeline:

         # The ratio of transformer1 to transformer2's output states to be combined during inference
         lowercase			:int							=  0.5

         # The shape of `encoder_hidden_states` is expected to be
         # `(batch_size, condition_lengths[0]+condition_lengths[1], num_features)`
         lowercase			:List[str]							=  [77, 2_57]

         # Which transformer to use to encode which condition.
         # E.g. `(1, 0)` means that we'll use `transformers[1](conditions[0])` and `transformers[0](conditions[1])`
         lowercase			:List[str]							=  [1, 0]






       def     SCREAMING_SNAKE_CASE			(      self:    Any				,    _lowerCAmelCase:    List[str]				,    _lowerCAmelCase:    Dict				,    _lowerCAmelCase:    str=None				,    _lowerCAmelCase:    Union[str, Any]=None				,    _lowerCAmelCase:    List[Any]=None				,    _lowerCAmelCase:    bool = True				,    ):
         lowercase			:str							=  hidden_states

         lowercase			:Dict							=  []
         lowercase			:Tuple							=  0
         # attention_mask is not used yet
         for i in range(2 ):
           # for each of the two transformers, pass the corresponding condition tokens
           lowercase			:int							=  encoder_hidden_states[:, tokens_start : tokens_start + self.condition_lengths[i]]
           lowercase			:Tuple							=  self.transformer_index_for_condition[i]
           lowercase			:int							=  self.transformers[transformer_index](
               _lowerCAmelCase				,    encoder_hidden_states=_lowerCAmelCase				,    timestep=_lowerCAmelCase				,    cross_attention_kwargs=_lowerCAmelCase				,    return_dict=_lowerCAmelCase				,    )[0]
           encoded_states.append(encoded_state - input_states )
           tokens_start += self.condition_lengths[i]

         lowercase			:int							=  encoded_states[0] * self.mix_ratio + encoded_states[1] * (1 - self.mix_ratio)
         lowercase			:Union[str, Any]							=  output_states + input_states

         if not return_dict:
           return (output_states,)

         return TransformeraDModelOutput(sample=_lowerCAmelCase )









_UpperCAmelCase		:  Tuple	=   {str(digit): digit**5 for digit in range(10)}



def 					UpperCAmelCase__     (   lowerCamelCase      ):
  return sum(DIGITS_FIFTH_POWER[digit] for digit in str(lowerCamelCase      )      )



def 					UpperCAmelCase__     (   ):
  return sum(
      number
      for number in range(1000,	1000000      )
      if number == digits_fifth_powers_sum(lowerCamelCase      )      )


if __name__ == "__main__":
     print(solution())



