


'''simple docstring'''
from __future__ import annotations

import unittest

from transformers import LEDConfig, is_tf_available
from transformers.testing_utils import require_tf, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
 import tensorflow as tf

 from transformers import TFLEDForConditionalGeneration, TFLEDModel

@require_tf
class     _lowerCamelCase       :




 '''simple docstring'''
 __lowercase					:		Union[str, Any]     				=    LEDConfig
 __lowercase					:		List[str]     				=    {}
 __lowercase					:		Optional[Any]     				=    """gelu"""


 def __init__(   self   ,   __lowercase   ,   __lowercase=13   ,   __lowercase=7   ,   __lowercase=True   ,   __lowercase=False   ,   __lowercase=99   ,   __lowercase=32   ,   __lowercase=2   ,   __lowercase=4   ,   __lowercase=37   ,   __lowercase=0.1   ,   __lowercase=0.1   ,   __lowercase=20   ,   __lowercase=2   ,   __lowercase=1   ,   __lowercase=0   ,   __lowercase=4   ,   ):

  """simple docstring"""







  __A	:		Optional[int]	  =      parent
  __A	:		Optional[int]	  =      batch_size
  __A	:		Dict	  =      seq_length
  __A	:		List[Any]	  =      is_training
  __A	:		List[str]	  =      use_labels
  __A	:		Dict	  =      vocab_size
  __A	:		Optional[Any]	  =      hidden_size
  __A	:		int	  =      num_hidden_layers
  __A	:		str	  =      num_attention_heads
  __A	:		Dict	  =      intermediate_size

  __A	:		Any	  =      hidden_dropout_prob
  __A	:		str	  =      attention_probs_dropout_prob
  __A	:		str	  =      max_position_embeddings
  __A	:		Optional[Any]	  =      eos_token_id
  __A	:		str	  =      pad_token_id
  __A	:		int	  =      bos_token_id
  __A	:		int	  =      attention_window

  # `ModelTesterMixin.test_attention_outputs` is expecting attention tensors to be of size
  # [num_attention_heads, encoder_seq_length, encoder_key_length], but TFLongformerSelfAttention
  # returns attention of shape [num_attention_heads, encoder_seq_length, self.attention_window + 1]
  # because its local attention only attends to `self.attention_window` and one before and one after
  __A	:		List[str]	  =      self.attention_window + 2

  # because of padding `encoder_seq_length`, is different from `seq_length`. Relevant for
  # the `test_attention_outputs` and `test_hidden_states_output` tests
  __A	:		Any	  =      (
      self.seq_length + (self.attention_window - self.seq_length % self.attention_window) % self.attention_window
  )


 def 			snake_case__ (   self    ):

  """simple docstring"""







  __A	:		List[Any]	  =      ids_tensor([self.batch_size, self.seq_length - 1]   ,   self.vocab_size    )
  __A	:		Any	  =      tf.expand_dims(tf.constant([self.eos_token_id] * self.batch_size    )   ,   1    )
  __A	:		Any	  =      tf.concat([input_ids, eos_tensor]   ,   axis=1    )

  __A	:		Tuple	  =      ids_tensor([self.batch_size, self.seq_length]   ,   self.vocab_size    )

  __A	:		str	  =      self.config_cls(
      vocab_size=self.vocab_size   ,   d_model=self.hidden_size   ,   encoder_layers=self.num_hidden_layers   ,   decoder_layers=self.num_hidden_layers   ,   encoder_attention_heads=self.num_attention_heads   ,   decoder_attention_heads=self.num_attention_heads   ,   encoder_ffn_dim=self.intermediate_size   ,   decoder_ffn_dim=self.intermediate_size   ,   dropout=self.hidden_dropout_prob   ,   attention_dropout=self.attention_probs_dropout_prob   ,   max_position_embeddings=self.max_position_embeddings   ,   eos_token_ids=[2]   ,   bos_token_id=self.bos_token_id   ,   pad_token_id=self.pad_token_id   ,   decoder_start_token_id=self.pad_token_id   ,   attention_window=self.attention_window   ,   **self.config_updates   ,   )
  __A	:		Union[str, Any]	  =      prepare_led_inputs_dict(__lowercase   ,   __lowercase   ,   __lowercase    )
  __A	:		Union[str, Any]	  =      tf.concat(
      [tf.zeros_like(__lowercase    )[:, :-1], tf.ones_like(__lowercase    )[:, -1:]]   ,   axis=-1   ,   )
  __A	:		Union[str, Any]	  =      global_attention_mask
  return config, inputs_dict







 def 			snake_case__ (   self   ,   __lowercase   ,   __lowercase    ):

  """simple docstring"""







  __A	:		Dict	  =      TFLEDModel(config=__lowercase    ).get_decoder()
  __A	:		int	  =      inputs_dict['input_ids']

  __A	:		Any	  =      input_ids[:1, :]
  __A	:		Optional[Any]	  =      inputs_dict['attention_mask'][:1, :]
  __A	:		List[Any]	  =      1

  # first forward pass
  __A	:		Tuple	  =      model(__lowercase   ,   attention_mask=__lowercase   ,   use_cache=__lowercase    )

  __A  ,__A	:		List[str]	  =      outputs.to_tuple()

  # create hypothetical next token and extent to next_input_ids
  __A	:		Optional[Any]	  =      ids_tensor((self.batch_size, 3)   ,   config.vocab_size    )
  __A	:		Optional[int]	  =      tf.cast(ids_tensor((self.batch_size, 3)   ,   2    )   ,   tf.inta    )

  # append to next input_ids and
  __A	:		Union[str, Any]	  =      tf.concat([input_ids, next_tokens]   ,   axis=-1    )
  __A	:		int	  =      tf.concat([attention_mask, next_attn_mask]   ,   axis=-1    )

  __A	:		Any	  =      model(__lowercase   ,   attention_mask=__lowercase    )[0]
  __A	:		Any	  =      model(__lowercase   ,   attention_mask=__lowercase   ,   past_key_values=__lowercase    )[0]

  self.parent.assertEqual(next_tokens.shape[1]   ,   output_from_past.shape[1]    )

  # select random slice
  __A	:		int	  =      int(ids_tensor((1,)   ,   output_from_past.shape[-1]    )    )
  __A	:		str	  =      output_from_no_past[:, -3:, random_slice_idx]
  __A	:		Dict	  =      output_from_past[:, :, random_slice_idx]

  # test that outputs are equal for slice
  tf.debugging.assert_near(__lowercase   ,   __lowercase   ,   rtol=1E-3    )


def       _lowercase							(					UpperCamelCase__							:       Any,					UpperCamelCase__							:       Optional[int],					UpperCamelCase__							:       Optional[int],					UpperCamelCase__							:       List[Any]=None,					UpperCamelCase__							:       Union[str, Any]=None,					UpperCamelCase__							:       Tuple=None,					UpperCamelCase__							:       Optional[int]=None,					):
 if attention_mask is None:
  __A	:		int	  =      tf.cast(tf.math.not_equal(UpperCamelCase__,					config.pad_token_id      ),					tf.inta      )
 if decoder_attention_mask is None:
  __A	:		Dict	  =      tf.concat(
      [
          tf.ones(decoder_input_ids[:, :1].shape,					dtype=tf.inta      ),
          tf.cast(tf.math.not_equal(decoder_input_ids[:, 1:],					config.pad_token_id      ),					tf.inta      ),
      ],					axis=-1,					)
 if head_mask is None:
  __A	:		Tuple	  =      tf.ones((config.encoder_layers, config.encoder_attention_heads)      )
 if decoder_head_mask is None:
  __A	:		str	  =      tf.ones((config.decoder_layers, config.decoder_attention_heads)      )
 return {
     "input_ids": input_ids,
     "attention_mask": attention_mask,
     "decoder_input_ids": decoder_input_ids,
     "decoder_attention_mask": decoder_attention_mask,
     "head_mask": head_mask,
     "decoder_head_mask": decoder_head_mask,
 }

@require_tf
class     _lowerCamelCase       (     snake_case_     ,		snake_case_     ,		unittest.TestCase					):




 '''simple docstring'''
 __lowercase					:		Union[str, Any]     				=    (TFLEDForConditionalGeneration, TFLEDModel) if is_tf_available() else ()
 __lowercase					:		List[str]     				=    (TFLEDForConditionalGeneration,) if is_tf_available() else ()
 __lowercase					:		int     				=    (
     {
         """conversational""": TFLEDForConditionalGeneration,
         """feature-extraction""": TFLEDModel,
         """summarization""": TFLEDForConditionalGeneration,
         """text2text-generation""": TFLEDForConditionalGeneration,
         """translation""": TFLEDForConditionalGeneration,
     }
     if is_tf_available()
     else {}
 )
 __lowercase					:		List[Any]     				=    True
 __lowercase					:		Dict     				=    False
 __lowercase					:		Union[str, Any]     				=    False
 __lowercase					:		Tuple     				=    False


 def 			snake_case__ (   self    ):

  """simple docstring"""







  __A	:		int	  =      TFLEDModelTester(self    )
  __A	:		Union[str, Any]	  =      ConfigTester(self   ,   config_class=__lowercase    )


 def 			snake_case__ (   self    ):

  """simple docstring"""







  self.config_tester.run_common_tests()


 def 			snake_case__ (   self    ):

  """simple docstring"""







  __A	:		str	  =      self.model_tester.prepare_config_and_inputs_for_common()
  self.model_tester.check_decoder_model_past_large_inputs(*__lowercase    )


 def 			snake_case__ (   self    ):

  """simple docstring"""







  __A  ,__A	:		Union[str, Any]	  =      self.model_tester.prepare_config_and_inputs_for_common()
  __A	:		List[str]	  =      tf.zeros_like(inputs_dict['attention_mask']    )
  __A	:		Optional[Any]	  =      2
  __A	:		Optional[int]	  =      tf.where(
      tf.range(self.model_tester.seq_length    )[None, :] < num_global_attn_indices   ,   1   ,   inputs_dict['global_attention_mask']   ,   )

  __A	:		List[Any]	  =      True
  __A	:		Union[str, Any]	  =      self.model_tester.seq_length
  __A	:		str	  =      self.model_tester.encoder_seq_length

  def check_decoder_attentions_output(__lowercase    ):
   __A	:		Tuple	  =      outputs.decoder_attentions
   self.assertEqual(len(__lowercase    )   ,   self.model_tester.num_hidden_layers    )
   self.assertListEqual(
       list(decoder_attentions[0].shape[-3:]    )   ,   [self.model_tester.num_attention_heads, seq_length, seq_length]   ,   )

  def check_encoder_attentions_output(__lowercase    ):
   __A	:		List[Any]	  =      [t.numpy() for t in outputs.encoder_attentions]
   __A	:		Optional[int]	  =      [t.numpy() for t in outputs.encoder_global_attentions]
   self.assertEqual(len(__lowercase    )   ,   self.model_tester.num_hidden_layers    )
   self.assertEqual(len(__lowercase    )   ,   self.model_tester.num_hidden_layers    )
   self.assertListEqual(
       list(attentions[0].shape[-3:]    )   ,   [self.model_tester.num_attention_heads, seq_length, seq_length]   ,   )
   self.assertListEqual(
       list(global_attentions[0].shape[-3:]    )   ,   [self.model_tester.num_attention_heads, encoder_seq_length, num_global_attn_indices]   ,   )

  for model_class in self.all_model_classes:
   __A	:		str	  =      True
   __A	:		str	  =      False
   __A	:		Union[str, Any]	  =      False
   __A	:		Dict	  =      model_class(__lowercase    )
   __A	:		Dict	  =      model(self._prepare_for_class(__lowercase   ,   __lowercase    )    )
   __A	:		Optional[int]	  =      len(__lowercase    )
   self.assertEqual(config.output_hidden_states   ,   __lowercase    )
   check_encoder_attentions_output(__lowercase    )

   if self.is_encoder_decoder:
    __A	:		Any	  =      model_class(__lowercase    )
    __A	:		str	  =      model(self._prepare_for_class(__lowercase   ,   __lowercase    )    )
    self.assertEqual(config.output_hidden_states   ,   __lowercase    )
    check_decoder_attentions_output(__lowercase    )

   # Check that output attentions can also be changed via the config
   del inputs_dict["output_attentions"]
   __A	:		List[str]	  =      True
   __A	:		Optional[Any]	  =      model_class(__lowercase    )
   __A	:		Tuple	  =      model(self._prepare_for_class(__lowercase   ,   __lowercase    )    )
   self.assertEqual(config.output_hidden_states   ,   __lowercase    )
   check_encoder_attentions_output(__lowercase    )

   # Check attention is always last and order is fine
   __A	:		str	  =      True
   __A	:		Dict	  =      True
   __A	:		List[Any]	  =      model_class(__lowercase    )
   __A	:		List[Any]	  =      model(self._prepare_for_class(__lowercase   ,   __lowercase    )    )

   self.assertEqual(out_len + (2 if self.is_encoder_decoder else 1)   ,   len(__lowercase    )    )
   self.assertEqual(model.config.output_hidden_states   ,   __lowercase    )
   check_encoder_attentions_output(__lowercase    )


 @unittest.skip('LED keeps using potentially symbolic tensors in conditionals and breaks tracing.'    )
 def 			snake_case__ (   self    ):

  """simple docstring"""







  pass







 def 			snake_case__ (   self    ):

  """simple docstring"""







  pass


def       _lowercase							(					UpperCamelCase__							:       int      ):
 return tf.constant(UpperCamelCase__,					dtype=tf.intaa      )


UpperCAmelCase_			:  Any										=   1E-4


@slow
@require_tf
class     _lowerCamelCase       (     unittest.TestCase					):




 '''simple docstring'''


 def 			snake_case__ (   self    ):

  """simple docstring"""







  __A	:		Union[str, Any]	  =      TFLEDForConditionalGeneration.from_pretrained('allenai/led-base-16384'    ).led

  # change to intended input here
  __A	:		int	  =      _long_tensor([512 * [0, 31_414, 232, 328, 740, 1_140, 12_695, 69]]    )
  __A	:		Union[str, Any]	  =      _long_tensor([128 * [0, 31_414, 232, 328, 740, 1_140, 12_695, 69]]    )
  __A	:		Optional[Any]	  =      prepare_led_inputs_dict(model.config   ,   __lowercase   ,   __lowercase    )
  __A	:		Optional[Any]	  =      model(**__lowercase    )[0]
  __A	:		Optional[int]	  =      (1, 1_024, 768)
  self.assertEqual(output.shape   ,   __lowercase    )
  # change to expected output here
  __A	:		Tuple	  =      tf.convert_to_tensor(
      [[2.3_0_5_0, 2.8_2_7_9, 0.6_5_3_1], [-1.8_4_5_7, -0.1_4_5_5, -3.5_6_6_1], [-1.0_1_8_6, 0.4_5_8_6, -2.2_0_4_3]]   ,   )
  tf.debugging.assert_near(output[:, :3, :3]   ,   __lowercase   ,   atol=1E-3    )







 def 			snake_case__ (   self    ):

  """simple docstring"""







  __A	:		List[Any]	  =      TFLEDForConditionalGeneration.from_pretrained('allenai/led-base-16384'    )

  # change to intended input here
  __A	:		Optional[int]	  =      _long_tensor([512 * [0, 31_414, 232, 328, 740, 1_140, 12_695, 69]]    )
  __A	:		int	  =      _long_tensor([128 * [0, 31_414, 232, 328, 740, 1_140, 12_695, 69]]    )
  __A	:		List[str]	  =      prepare_led_inputs_dict(model.config   ,   __lowercase   ,   __lowercase    )
  __A	:		Any	  =      model(**__lowercase    )[0]
  __A	:		Optional[int]	  =      (1, 1_024, model.config.vocab_size)
  self.assertEqual(output.shape   ,   __lowercase    )
  # change to expected output here
  __A	:		str	  =      tf.convert_to_tensor(
      [[3_3.6_5_0_7, 6.4_5_7_2, 1_6.8_0_8_9], [5.8_7_3_9, -2.4_2_3_8, 1_1.2_9_0_2], [-3.2_1_3_9, -4.3_1_4_9, 4.2_7_8_3]]   ,   )
  tf.debugging.assert_near(output[:, :3, :3]   ,   __lowercase   ,   atol=1E-3   ,   rtol=1E-3    )







"""simple docstring"""



import os
from typing import Any, Callable, Dict, List, Optional, Tuple, Union

import torch
from torch import nn

from ...models.controlnet import ControlNetModel, ControlNetOutput
from ...models.modeling_utils import ModelMixin
from ...utils import logging


SCREAMING_SNAKE_CASE__:Optional[Any]									=			logging.get_logger(__name__)






class  snake_case__      (       snake_case_       ):


 def __init__( self    ,		lowerCamelCase						):
        super().__init__()
        __a       =    nn.ModuleList(lowerCamelCase						)


 def 						a__   ( self    ,		lowerCamelCase    ,		lowerCamelCase    ,		lowerCamelCase    ,		lowerCamelCase    ,		lowerCamelCase    ,		lowerCamelCase = None    ,		lowerCamelCase = None    ,		lowerCamelCase = None    ,		lowerCamelCase = None    ,		lowerCamelCase = False    ,		lowerCamelCase = True    ,		):
        for i, (image, scale, controlnet) in enumerate(zip(lowerCamelCase    ,		lowerCamelCase    ,		self.nets						)						):
               __a ,   __a       =    controlnet(
                   lowerCamelCase    ,		lowerCamelCase    ,		lowerCamelCase    ,		lowerCamelCase    ,		lowerCamelCase    ,		lowerCamelCase    ,		lowerCamelCase    ,		lowerCamelCase    ,		lowerCamelCase    ,		lowerCamelCase    ,		lowerCamelCase    ,		)

               # merge samples
               if i == 0:
                      __a ,   __a       =    down_samples, mid_sample
               else:
                      __a       =    [
                          samples_prev + samples_curr
                          for samples_prev, samples_curr in zip(lowerCamelCase    ,		lowerCamelCase						)
                      ]
                      mid_block_res_sample += mid_sample

        return down_block_res_samples, mid_block_res_sample


 def 						a__   ( self    ,		lowerCamelCase    ,		lowerCamelCase = True    ,		lowerCamelCase = None    ,		lowerCamelCase = False    ,		lowerCamelCase = None    ,		):
        __a       =    0
        __a       =    save_directory
        for controlnet in self.nets:
               controlnet.save_pretrained(
                   lowerCamelCase    ,		is_main_process=lowerCamelCase    ,		save_function=lowerCamelCase    ,		safe_serialization=lowerCamelCase    ,		variant=lowerCamelCase    ,		)

               idx += 1
               __a       =    model_path_to_save + F"_{idx}"






 @classmethod
 def 						a__   ( cls    ,		lowerCamelCase    ,		**lowerCamelCase						):
        __a       =    0
        __a       =    []

        # load controlnet and append to list until no controlnet directory exists anymore
        # first controlnet has to be saved under `./mydirectory/controlnet` to be compliant with `DiffusionPipeline.from_prertained`
        # second, third, ... controlnets have to be saved under `./mydirectory/controlnet_1`, `./mydirectory/controlnet_2`, ...
        __a       =    pretrained_model_path
        while os.path.isdir(lowerCamelCase						):
               __a       =    ControlNetModel.from_pretrained(lowerCamelCase    ,		**lowerCamelCase						)
               controlnets.append(lowerCamelCase						)

               idx += 1
               __a       =    pretrained_model_path + F"_{idx}"

        logger.info(F"{len(lowerCamelCase						)} controlnets loaded from {pretrained_model_path}."						)

        if len(lowerCamelCase						) == 0:
               raise ValueError(
                   F"No ControlNets found under {os.path.dirname(lowerCamelCase						)}. Expected at least {pretrained_model_path + '_0'}."						)

        return cls(lowerCamelCase						)
