'''simple docstring'''




from __future__ import annotations

import unittest

from transformers import 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, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
  import numpy
  import tensorflow as tf

  from transformers import (
      TF_DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST,
      TF_DPR_QUESTION_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST,
      TF_DPR_READER_PRETRAINED_MODEL_ARCHIVE_LIST,
      BertConfig,
      DPRConfig,
      TFDPRContextEncoder,
      TFDPRQuestionEncoder,
      TFDPRReader,
  )

class 	lowerCamelCase_	:
      '''simple docstring'''

      def __init__(       self						:				int						,					A						:				Any						,					A						:				Dict=13						,					A						:				List[Any]=7						,					A						:				Dict=True						,					A						:				Dict=True						,					A						:				Union[str, Any]=True						,					A						:				Dict=True						,					A						:				Optional[int]=99						,					A						:				Union[str, Any]=32						,					A						:				List[str]=2						,					A						:				Dict=4						,					A						:				List[Any]=37						,					A						:				str="gelu"						,					A						:				List[Any]=0.1						,					A						:				int=0.1						,					A						:				Optional[Any]=512						,					A						:				str=16						,					A						:				Dict=2						,					A						:				Dict=0.02						,					A						:				Optional[int]=3						,					A						:				int=4						,					A						:				int=None						,					A						:				Tuple=0						,					):
             _UpperCAmelCase     :  int							      =    parent
             _UpperCAmelCase     :  str							      =    batch_size
             _UpperCAmelCase     :  List[str]							      =    seq_length
             _UpperCAmelCase     :  Optional[int]							      =    is_training
             _UpperCAmelCase     :  Dict							      =    use_input_mask
             _UpperCAmelCase     :  str							      =    use_token_type_ids
             _UpperCAmelCase     :  List[Any]							      =    use_labels
             _UpperCAmelCase     :  int							      =    vocab_size
             _UpperCAmelCase     :  Optional[int]							      =    hidden_size
             _UpperCAmelCase     :  Dict							      =    num_hidden_layers
             _UpperCAmelCase     :  Optional[int]							      =    num_attention_heads
             _UpperCAmelCase     :  Union[str, Any]							      =    intermediate_size
             _UpperCAmelCase     :  Union[str, Any]							      =    hidden_act
             _UpperCAmelCase     :  List[str]							      =    hidden_dropout_prob
             _UpperCAmelCase     :  int							      =    attention_probs_dropout_prob
             _UpperCAmelCase     :  List[Any]							      =    max_position_embeddings
             _UpperCAmelCase     :  List[Any]							      =    type_vocab_size
             _UpperCAmelCase     :  Union[str, Any]							      =    type_sequence_label_size
             _UpperCAmelCase     :  Optional[int]							      =    initializer_range
             _UpperCAmelCase     :  str							      =    num_labels
             _UpperCAmelCase     :  int							      =    num_choices
             _UpperCAmelCase     :  List[Any]							      =    scope
             _UpperCAmelCase     :  str							      =    projection_dim




      def      _A	(       self						:				Tuple       ):
             _UpperCAmelCase     :  Union[str, Any]							      =    ids_tensor([self.batch_size, self.seq_length]						,					self.vocab_size       )

             _UpperCAmelCase     :  Optional[int]							      =    None
             if self.use_input_mask:
                    # follow test_modeling_tf_ctrl.py
                    _UpperCAmelCase     :  List[str]							      =    random_attention_mask([self.batch_size, self.seq_length]       )

             _UpperCAmelCase     :  str							      =    None
             if self.use_token_type_ids:
                    _UpperCAmelCase     :  str							      =    ids_tensor([self.batch_size, self.seq_length]						,					self.type_vocab_size       )

             _UpperCAmelCase     :  Optional[Any]							      =    None
             _UpperCAmelCase     :  List[str]							      =    None
             _UpperCAmelCase     :  Union[str, Any]							      =    None
             if self.use_labels:
                    _UpperCAmelCase     :  int							      =    ids_tensor([self.batch_size]						,					self.type_sequence_label_size       )
                    _UpperCAmelCase     :  Optional[Any]							      =    ids_tensor([self.batch_size, self.seq_length]						,					self.num_labels       )
                    _UpperCAmelCase     :  Any							      =    ids_tensor([self.batch_size]						,					self.num_choices       )

             _UpperCAmelCase     :  Optional[int]							      =    BertConfig(
                 vocab_size=self.vocab_size						,					hidden_size=self.hidden_size						,					num_hidden_layers=self.num_hidden_layers						,					num_attention_heads=self.num_attention_heads						,					intermediate_size=self.intermediate_size						,					hidden_act=self.hidden_act						,					hidden_dropout_prob=self.hidden_dropout_prob						,					attention_probs_dropout_prob=self.attention_probs_dropout_prob						,					max_position_embeddings=self.max_position_embeddings						,					type_vocab_size=self.type_vocab_size						,					is_decoder=A						,					initializer_range=self.initializer_range						,					)
             _UpperCAmelCase     :  List[str]							      =    DPRConfig(projection_dim=self.projection_dim						,					**config.to_dict()       )

             return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels




      def      _A	(       self						:				int						,					A						:				str						,					A						:				Optional[Any]						,					A						:				str						,					A						:				List[Any]						,					A						:				Optional[int]						,					A						:				Optional[int]						,					A						:				int       ):
             _UpperCAmelCase     :  Optional[Any]							      =    TFDPRContextEncoder(config=A       )
             _UpperCAmelCase     :  Optional[int]							      =    model(A						,					attention_mask=A						,					token_type_ids=A       )
             _UpperCAmelCase     :  Optional[Any]							      =    model(A						,					token_type_ids=A       )
             _UpperCAmelCase     :  List[str]							      =    model(A       )
             self.parent.assertEqual(result.pooler_output.shape						,					(self.batch_size, self.projection_dim or self.hidden_size)       )




      def      _A	(       self						:				str						,					A						:				int						,					A						:				Optional[int]						,					A						:				int						,					A						:				List[Any]						,					A						:				int						,					A						:				Union[str, Any]						,					A						:				List[Any]       ):
             _UpperCAmelCase     :  Optional[Any]							      =    TFDPRQuestionEncoder(config=A       )
             _UpperCAmelCase     :  List[Any]							      =    model(A						,					attention_mask=A						,					token_type_ids=A       )
             _UpperCAmelCase     :  str							      =    model(A						,					token_type_ids=A       )
             _UpperCAmelCase     :  List[str]							      =    model(A       )
             self.parent.assertEqual(result.pooler_output.shape						,					(self.batch_size, self.projection_dim or self.hidden_size)       )




      def      _A	(       self						:				Optional[int]						,					A						:				Any						,					A						:				Dict						,					A						:				Union[str, Any]						,					A						:				Any						,					A						:				List[str]						,					A						:				Tuple						,					A						:				List[Any]       ):
             _UpperCAmelCase     :  int							      =    TFDPRReader(config=A       )
             _UpperCAmelCase     :  int							      =    model(A						,					attention_mask=A       )

             self.parent.assertEqual(result.start_logits.shape						,					(self.batch_size, self.seq_length)       )
             self.parent.assertEqual(result.end_logits.shape						,					(self.batch_size, self.seq_length)       )
             self.parent.assertEqual(result.relevance_logits.shape						,					(self.batch_size,)       )





      def      _A	(       self						:				Union[str, Any]       ):
             _UpperCAmelCase     :  str							      =    self.prepare_config_and_inputs()
             (
                 (
                 _UpperCAmelCase
             )      ,  (
                 _UpperCAmelCase
             )      ,  (
                 _UpperCAmelCase
             )      ,  (
                 _UpperCAmelCase
             )      ,  (
                 _UpperCAmelCase
             )      ,  (
                 _UpperCAmelCase
             )      ,  (
                 _UpperCAmelCase
             )      ,  
             )     :  List[Any]							      =    config_and_inputs
             _UpperCAmelCase     :  Dict							      =    {"input_ids": input_ids}
             return config, inputs_dict

@require_tf
class 	lowerCamelCase_	(snake_case__							, snake_case__							, unittest.TestCase    ):
      '''simple docstring'''

      __UpperCamelCase:  Union[str, Any]						=						(
          (
              TFDPRContextEncoder,
              TFDPRQuestionEncoder,
              TFDPRReader,
          )
          if is_tf_available()
          else ()
      )
      __UpperCamelCase:  int						=						{"feature-extraction": TFDPRQuestionEncoder} if is_tf_available() else {}

      __UpperCamelCase:  Optional[int]						=						False
      __UpperCamelCase:  Dict						=						False
      __UpperCamelCase:  int						=						False
      __UpperCamelCase:  str						=						False
      __UpperCamelCase:  int						=						False




      def      _A	(       self						:				Optional[Any]       ):
             _UpperCAmelCase     :  Tuple							      =    TFDPRModelTester(self       )
             _UpperCAmelCase     :  Optional[Any]							      =    ConfigTester(self						,					config_class=A						,					hidden_size=37       )




      def      _A	(       self						:				str       ):
             self.config_tester.run_common_tests()




      def      _A	(       self						:				int       ):
             _UpperCAmelCase     :  Dict							      =    self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_dpr_context_encoder(*A       )




      def      _A	(       self						:				int       ):
             _UpperCAmelCase     :  Optional[int]							      =    self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_dpr_question_encoder(*A       )




      def      _A	(       self						:				List[str]       ):
             _UpperCAmelCase     :  Tuple							      =    self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_dpr_reader(*A       )





      @slow
      def      _A	(       self						:				Optional[Any]       ):
             for model_name in TF_DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                    _UpperCAmelCase     :  int							      =    TFDPRContextEncoder.from_pretrained(A       )
                    self.assertIsNotNone(A       )

             for model_name in TF_DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                    _UpperCAmelCase     :  List[Any]							      =    TFDPRContextEncoder.from_pretrained(A       )
                    self.assertIsNotNone(A       )

             for model_name in TF_DPR_QUESTION_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                    _UpperCAmelCase     :  List[str]							      =    TFDPRQuestionEncoder.from_pretrained(A       )
                    self.assertIsNotNone(A       )

             for model_name in TF_DPR_READER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                    _UpperCAmelCase     :  Dict							      =    TFDPRReader.from_pretrained(A       )
                    self.assertIsNotNone(A       )

@require_tf
class 	lowerCamelCase_	(unittest.TestCase    ):
      '''simple docstring'''

      @slow
      def      _A	(       self						:				Any       ):
             _UpperCAmelCase     :  str							      =    TFDPRQuestionEncoder.from_pretrained("facebook/dpr-question_encoder-single-nq-base"       )

             _UpperCAmelCase     :  List[str]							      =    tf.constant(
                 [[101, 7592, 1010, 2003, 2026, 3899, 10140, 1029, 102]]       )  # [CLS] hello, is my dog cute? [SEP]
             _UpperCAmelCase     :  List[str]							      =    model(A       )[0]  # embedding shape = (1, 768)
             # compare the actual values for a slice.
             _UpperCAmelCase     :  List[Any]							      =    tf.constant(
                 [
                     [
                         0.03_236_253,
                         0.12_753_335,
                         0.16_818_509,
                         0.00_279_786,
                         0.3_896_933,
                         0.24_264_945,
                         0.2_178_971,
                         -0.02_335_227,
                         -0.08_481_959,
                         -0.14_324_117,
                     ]
                 ]       )
             self.assertTrue(numpy.allclose(output[:, :10].numpy()						,					expected_slice.numpy()						,					atol=1E-4       )       )





'''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


a_  :				Dict						= logging.get_logger(__name__)





class      snake_case    (    lowercase			):





  """simple docstring"""




  def __init__(    self			,  UpperCamelCase     ):







   """simple docstring"""
   super().__init__()
   lowerCamelCase_      =  nn.ModuleList(UpperCamelCase     )






  def   snake_case		(    self			,  UpperCamelCase			,  UpperCamelCase			,  UpperCamelCase			,  UpperCamelCase			,  UpperCamelCase			,  UpperCamelCase = None			,  UpperCamelCase = None			,  UpperCamelCase = None			,  UpperCamelCase = None			,  UpperCamelCase = False			,  UpperCamelCase = True			,  ):







   """simple docstring"""
   for i, (image, scale, controlnet) in enumerate(zip(UpperCamelCase			,  UpperCamelCase			,  self.nets     )     ):
    lowerCamelCase_ ,lowerCamelCase_      =  controlnet(
        UpperCamelCase			,  UpperCamelCase			,  UpperCamelCase			,  UpperCamelCase			,  UpperCamelCase			,  UpperCamelCase			,  UpperCamelCase			,  UpperCamelCase			,  UpperCamelCase			,  UpperCamelCase			,  UpperCamelCase			,  )

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

   return down_block_res_samples, mid_block_res_sample






  def   snake_case		(    self			,  UpperCamelCase			,  UpperCamelCase = True			,  UpperCamelCase = None			,  UpperCamelCase = False			,  UpperCamelCase = None			,  ):







   """simple docstring"""
   lowerCamelCase_      =  0
   lowerCamelCase_      =  save_directory
   for controlnet in self.nets:
    controlnet.save_pretrained(
        UpperCamelCase			,  is_main_process=UpperCamelCase			,  save_function=UpperCamelCase			,  safe_serialization=UpperCamelCase			,  variant=UpperCamelCase			,  )

    idx += 1
    lowerCamelCase_      =  model_path_to_save + f'''_{idx}'''







  @classmethod
  def   snake_case		(    cls			,  UpperCamelCase			,  **UpperCamelCase     ):







   """simple docstring"""
   lowerCamelCase_      =  0
   lowerCamelCase_      =  []

   # 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`, ...
   lowerCamelCase_      =  pretrained_model_path
   while os.path.isdir(UpperCamelCase     ):
    lowerCamelCase_      =  ControlNetModel.from_pretrained(UpperCamelCase			,  **UpperCamelCase     )
    controlnets.append(UpperCamelCase     )

    idx += 1
    lowerCamelCase_      =  pretrained_model_path + f'''_{idx}'''

   logger.info(f'''{len(UpperCamelCase     )} controlnets loaded from {pretrained_model_path}.'''     )

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

   return cls(UpperCamelCase     )



