






'''simple docstring'''






from typing import Optional

from torch import nn

from .transformer_ad import TransformeraDModel, TransformeraDModelOutput






class  __magic_name__							(		nn.Module    ):






       def __init__(      self  ,							snake_case = 1_6  ,							snake_case = 8_8  ,							snake_case = None  ,							snake_case = 1  ,							snake_case = 0.0  ,							snake_case = 3_2  ,							snake_case = None  ,							snake_case = False  ,							snake_case = None  ,							snake_case = None  ,							snake_case = "geglu"  ,							snake_case = None  ,							)						->  List[str]:




          '''simple docstring'''
          super().__init__()
          _UpperCAmelCase       :       Any							   =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
          _UpperCAmelCase       :       Dict							   =0.5

          # The shape of `encoder_hidden_states` is expected to be
          # `(batch_size, condition_lengths[0]+condition_lengths[1], num_features)`
          _UpperCAmelCase       :       Optional[int]							   =[7_7, 2_5_7]

          # 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])`
          _UpperCAmelCase       :       List[Any]							   =[1, 0]






       def       lowerCAmelCase		(      self  ,							snake_case  ,							snake_case  ,							snake_case=None  ,							snake_case=None  ,							snake_case=None  ,							snake_case = True  ,							)						->  Tuple:




          '''simple docstring'''
          _UpperCAmelCase       :       Union[str, Any]							   =hidden_states

          _UpperCAmelCase       :       List[Any]							   =[]
          _UpperCAmelCase       :       List[Any]							   =0
          # attention_mask is not used yet
          for i in range(2):
             # for each of the two transformers, pass the corresponding condition tokens
             _UpperCAmelCase       :       Optional[Any]							   =encoder_hidden_states[:, tokens_start : tokens_start + self.condition_lengths[i]]
             _UpperCAmelCase       :       Union[str, Any]							   =self.transformer_index_for_condition[i]
             _UpperCAmelCase       :       Union[str, Any]							   =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]

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

          if not return_dict:
             return (output_states,)

          return TransformeraDModelOutput(sample=__lowerCAmelCase)





'''simple docstring'''






import shutil
import tempfile
import unittest

import numpy as np
import pytest

from transformers import is_speech_available, is_vision_available
from transformers.testing_utils import require_torch


if is_vision_available():
 from transformers import TvltImageProcessor

if is_speech_available():
 from transformers import TvltFeatureExtractor

from transformers import TvltProcessor






@require_torch
class  __magic_name__							(		unittest.TestCase    ):






       def       lowerCAmelCase		(      self)						->  str:




          '''simple docstring'''
          _UpperCAmelCase       :       Dict							   ='ZinengTang/tvlt-base'
          _UpperCAmelCase       :       Dict							   =tempfile.mkdtemp()






       def       lowerCAmelCase		(      self  ,							**snake_case)						->  Union[str, Any]:




          '''simple docstring'''
          return TvltImageProcessor.from_pretrained(self.checkpoint  ,							**snake_case)






       def       lowerCAmelCase		(      self  ,							**snake_case)						->  Dict:




          '''simple docstring'''
          return TvltFeatureExtractor.from_pretrained(self.checkpoint  ,							**snake_case)






       def       lowerCAmelCase		(      self)						->  Any:




          '''simple docstring'''
          shutil.rmtree(self.tmpdirname)






       def       lowerCAmelCase		(      self)						->  Optional[int]:




          '''simple docstring'''
          _UpperCAmelCase       :       Any							   =self.get_image_processor()
          _UpperCAmelCase       :       Optional[Any]							   =self.get_feature_extractor()

          _UpperCAmelCase       :       str							   =TvltProcessor(image_processor=snake_case  ,							feature_extractor=snake_case)
          processor.save_pretrained(self.tmpdirname)
          _UpperCAmelCase       :       str							   =TvltProcessor.from_pretrained(self.tmpdirname)

          self.assertIsInstance(processor.feature_extractor  ,							snake_case)
          self.assertIsInstance(processor.image_processor  ,							snake_case)






       def       lowerCAmelCase		(      self)						->  Union[str, Any]:




          '''simple docstring'''
          _UpperCAmelCase       :       str							   =self.get_image_processor()
          _UpperCAmelCase       :       List[Any]							   =self.get_feature_extractor()

          _UpperCAmelCase       :       str							   =TvltProcessor(image_processor=snake_case  ,							feature_extractor=snake_case)

          _UpperCAmelCase       :       Optional[int]							   =np.ones([1_2_0_0_0])

          _UpperCAmelCase       :       str							   =feature_extractor(snake_case  ,							return_tensors='np')
          _UpperCAmelCase       :       Dict							   =processor(audio=snake_case  ,							return_tensors='np')

          for key in audio_dict.keys():
             self.assertAlmostEqual(audio_dict[key].sum()  ,							input_processor[key].sum()  ,							delta=1E-2)






       def       lowerCAmelCase		(      self)						->  int:




          '''simple docstring'''
          _UpperCAmelCase       :       Dict							   =self.get_image_processor()
          _UpperCAmelCase       :       int							   =self.get_feature_extractor()

          _UpperCAmelCase       :       int							   =TvltProcessor(image_processor=snake_case  ,							feature_extractor=snake_case)

          _UpperCAmelCase       :       Union[str, Any]							   =np.ones([3, 2_2_4, 2_2_4])

          _UpperCAmelCase       :       Optional[Any]							   =image_processor(snake_case  ,							return_tensors='np')
          _UpperCAmelCase       :       List[Any]							   =processor(images=snake_case  ,							return_tensors='np')

          for key in image_dict.keys():
             self.assertAlmostEqual(image_dict[key].sum()  ,							input_processor[key].sum()  ,							delta=1E-2)






       def       lowerCAmelCase		(      self)						->  List[str]:




          '''simple docstring'''
          _UpperCAmelCase       :       List[str]							   =self.get_image_processor()
          _UpperCAmelCase       :       Dict							   =self.get_feature_extractor()

          _UpperCAmelCase       :       Optional[int]							   =TvltProcessor(image_processor=snake_case  ,							feature_extractor=snake_case)

          _UpperCAmelCase       :       Optional[int]							   =np.ones([1_2_0_0_0])
          _UpperCAmelCase       :       str							   =np.ones([3, 2_2_4, 2_2_4])

          _UpperCAmelCase       :       Optional[int]							   =processor(audio=snake_case  ,							images=snake_case)

          self.assertListEqual(list(inputs.keys())  ,							['audio_values', 'audio_mask', 'pixel_values', 'pixel_mask'])

          # test if it raises when no input is passed
          with pytest.raises(snake_case):
             processor()






       def       lowerCAmelCase		(      self)						->  Optional[int]:




          '''simple docstring'''
          _UpperCAmelCase       :       Optional[int]							   =self.get_image_processor()
          _UpperCAmelCase       :       Tuple							   =self.get_feature_extractor()

          _UpperCAmelCase       :       Dict							   =TvltProcessor(image_processor=snake_case  ,							feature_extractor=snake_case)

          self.assertListEqual(
              processor.model_input_names  ,							image_processor.model_input_names + feature_extractor.model_input_names  ,							msg='`processor` and `image_processor`+`feature_extractor` model input names do not match'  ,							)

