from typing import Optional

from torch import nn

from .transformer_ad import TransformeraDModel, TransformeraDModelOutput







class __snake_case    (						nn.Module		):




       def __init__(   self      :    Any	,     _lowercase      :    Any = 16	,     _lowercase      :    List[Any] = 88	,     _lowercase      :    int = None	,     _lowercase      :    Optional[Any] = 1	,     _lowercase      :    Dict = 0.0	,     _lowercase      :    List[Any] = 32	,     _lowercase      :    Tuple = None	,     _lowercase      :    List[str] = False	,     _lowercase      :    List[str] = None	,     _lowercase      :    List[Any] = None	,     _lowercase      :    Union[str, Any] = "geglu"	,     _lowercase      :    List[str] = None	,     ):



         """simple docstring"""






         super().__init__()
         SCREAMING_SNAKE_CASE__      							=				nn.ModuleList(
             [
                 TransformeraDModel(
                     num_attention_heads=_SCREAMING_SNAKE_CASE	,     attention_head_dim=_SCREAMING_SNAKE_CASE	,     in_channels=_SCREAMING_SNAKE_CASE	,     num_layers=_SCREAMING_SNAKE_CASE	,     dropout=_SCREAMING_SNAKE_CASE	,     norm_num_groups=_SCREAMING_SNAKE_CASE	,     cross_attention_dim=_SCREAMING_SNAKE_CASE	,     attention_bias=_SCREAMING_SNAKE_CASE	,     sample_size=_SCREAMING_SNAKE_CASE	,     num_vector_embeds=_SCREAMING_SNAKE_CASE	,     activation_fn=_SCREAMING_SNAKE_CASE	,     num_embeds_ada_norm=_SCREAMING_SNAKE_CASE	,     )
                 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
         SCREAMING_SNAKE_CASE__      							=				0.5

         # The shape of `encoder_hidden_states` is expected to be
         # `(batch_size, condition_lengths[0]+condition_lengths[1], num_features)`
         SCREAMING_SNAKE_CASE__      							=				[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])`
         SCREAMING_SNAKE_CASE__      							=				[1, 0]




       def      __a						(   self      :    Optional[Any]	,     _lowercase      :    Dict	,     _lowercase      :    List[Any]	,     _lowercase      :    List[str]=None	,     _lowercase      :    Union[str, Any]=None	,     _lowercase      :    str=None	,     _lowercase      :    int = True	,     ):



         """simple docstring"""






         SCREAMING_SNAKE_CASE__      							=				hidden_states

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

         SCREAMING_SNAKE_CASE__      							=				encoded_states[0] * self.mix_ratio + encoded_states[1] * (1 - self.mix_ratio)
         SCREAMING_SNAKE_CASE__      							=				output_states + input_states

         if not return_dict:
           return (output_states,)

         return TransformeraDModelOutput(sample=_SCREAMING_SNAKE_CASE				)

import contextlib
import importlib
import io
import unittest

import transformers

# Try to import everything from transformers to ensure every object can be loaded.
from transformers import *  # noqa F406
from transformers.testing_utils import DUMMY_UNKNOWN_IDENTIFIER, require_flax, require_tf, require_torch
from transformers.utils import ContextManagers, find_labels, is_flax_available, is_tf_available, is_torch_available


if is_torch_available():
   from transformers import BertForPreTraining, BertForQuestionAnswering, BertForSequenceClassification

if is_tf_available():
   from transformers import TFBertForPreTraining, TFBertForQuestionAnswering, TFBertForSequenceClassification

if is_flax_available():
   from transformers import FlaxBertForPreTraining, FlaxBertForQuestionAnswering, FlaxBertForSequenceClassification


__lowerCamelCase      :			Tuple							       =		DUMMY_UNKNOWN_IDENTIFIER
# An actual model hosted on huggingface.co

__lowerCamelCase      :			Optional[Any]							       =		'''main'''
# Default branch name
__lowerCamelCase      :			Optional[int]							       =		'''f2c752cfc5c0ab6f4bdec59acea69eefbee381c2'''
# One particular commit (not the top of `main`)
__lowerCamelCase      :			Any							       =		'''aaaaaaa'''
# This commit does not exist, so we should 404.

__lowerCamelCase      :			List[str]							       =		'''d9e9f15bc825e4b2c9249e9578f884bbcb5e3684'''
# Sha-1 of config.json on the top of `main`, for checking purposes
__lowerCamelCase      :			List[Any]							       =		'''4b243c475af8d0a7754e87d7d096c92e5199ec2fe168a2ee7998e3b8e9bcb1d3'''



@contextlib.contextmanager
def      __SCREAMING_SNAKE_CASE  ( )							-> str:







  """simple docstring"""





  print("""Welcome!"""     )
  yield
  print("""Bye!"""     )



@contextlib.contextmanager
def      __SCREAMING_SNAKE_CASE  ( )							-> Any:







  """simple docstring"""





  print("""Bonjour!"""     )
  yield
  print("""Au revoir!"""     )







class __snake_case    (						unittest.TestCase		):




   def      __a						(   self      :    List[Any]				):



     """simple docstring"""






     assert transformers.__spec__ is not None
     assert importlib.util.find_spec("""transformers"""				) is not None







class __snake_case    (						unittest.TestCase		):




   @unittest.mock.patch("""sys.stdout"""	,     new_callable=io.StringIO				)
   def      __a						(   self      :    List[str]	,     _lowercase      :    str				):



     """simple docstring"""






     with ContextManagers([]				):
       print("""Transformers are awesome!"""				)
     # The print statement adds a new line at the end of the output
     self.assertEqual(mock_stdout.getvalue()	,     """Transformers are awesome!\n"""				)




   @unittest.mock.patch("""sys.stdout"""	,     new_callable=io.StringIO				)
   def      __a						(   self      :    Optional[Any]	,     _lowercase      :    str				):



     """simple docstring"""






     with ContextManagers([context_en()]				):
       print("""Transformers are awesome!"""				)
     # The output should be wrapped with an English welcome and goodbye
     self.assertEqual(mock_stdout.getvalue()	,     """Welcome!\nTransformers are awesome!\nBye!\n"""				)




   @unittest.mock.patch("""sys.stdout"""	,     new_callable=io.StringIO				)
   def      __a						(   self      :    Tuple	,     _lowercase      :    Dict				):



     """simple docstring"""






     with ContextManagers([context_fr(), context_en()]				):
       print("""Transformers are awesome!"""				)
     # The output should be wrapped with an English and French welcome and goodbye
     self.assertEqual(mock_stdout.getvalue()	,     """Bonjour!\nWelcome!\nTransformers are awesome!\nBye!\nAu revoir!\n"""				)




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



     """simple docstring"""






     self.assertEqual(find_labels(_lowercase				)	,     ["""labels"""]				)
     self.assertEqual(find_labels(_lowercase				)	,     ["""labels""", """next_sentence_label"""]				)
     self.assertEqual(find_labels(_lowercase				)	,     ["""start_positions""", """end_positions"""]				)







     class __snake_case    (						lowerCamelCase_		):
        pass





     self.assertEqual(find_labels(_lowercase				)	,     ["""labels"""]				)




   @require_tf
   def      __a						(   self      :    Any				):



     """simple docstring"""






     self.assertEqual(find_labels(_lowercase				)	,     ["""labels"""]				)
     self.assertEqual(find_labels(_lowercase				)	,     ["""labels""", """next_sentence_label"""]				)
     self.assertEqual(find_labels(_lowercase				)	,     ["""start_positions""", """end_positions"""]				)







     class __snake_case    (						lowerCamelCase_		):
        pass





     self.assertEqual(find_labels(_lowercase				)	,     ["""labels"""]				)




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



     """simple docstring"""






     self.assertEqual(find_labels(_lowercase				)	,     []				)
     self.assertEqual(find_labels(_lowercase				)	,     []				)
     self.assertEqual(find_labels(_lowercase				)	,     []				)







     class __snake_case    (						lowerCamelCase_		):
        pass





     self.assertEqual(find_labels(_lowercase				)	,     []				)
