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


__UpperCAmelCase					=logging.get_logger(__name__)




class        a__ (  UpperCAmelCase__					):






     def __init__(					self						:						int    ,							a						:						Union[List[ControlNetModel], Tuple[ControlNetModel]]       ):





         """simple docstring"""



         super().__init__()
         __lowerCamelCase            =    nn.ModuleList(a       )






     def   SCREAMING_SNAKE_CASE__  (					self						:						Union[str, Any]    ,							a						:						torch.FloatTensor    ,							a						:						Union[torch.Tensor, float, int]    ,							a						:						torch.Tensor    ,							a						:						List[torch.tensor]    ,							a						:						List[float]    ,							a						:						Optional[torch.Tensor] = None    ,							a						:						Optional[torch.Tensor] = None    ,							a						:						Optional[torch.Tensor] = None    ,							a						:						Optional[Dict[str, Any]] = None    ,							a						:						bool = False    ,							a						:						bool = True    ,							):





         """simple docstring"""



         for i, (image, scale, controlnet) in enumerate(zip(a    ,							a    ,							self.nets       )       ):
             __lowerCamelCase      ,    __lowerCamelCase            =    controlnet(
                 a    ,							a    ,							a    ,							a    ,							a    ,							a    ,							a    ,							a    ,							a    ,							a    ,							a    ,							)

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

         return down_block_res_samples, mid_block_res_sample






     def   SCREAMING_SNAKE_CASE__  (					self						:						Dict    ,							a						:						Union[str, os.PathLike]    ,							a						:						bool = True    ,							a						:						Callable = None    ,							a						:						bool = False    ,							a						:						Optional[str] = None    ,							):





         """simple docstring"""



         __lowerCamelCase            =    0
         __lowerCamelCase            =    save_directory
         for controlnet in self.nets:
             controlnet.save_pretrained(
                 a    ,							is_main_process=a    ,							save_function=a    ,							safe_serialization=a    ,							variant=a    ,							)

             idx += 1
             __lowerCamelCase            =    model_path_to_save + f"""_{idx}"""






     @classmethod
     def   SCREAMING_SNAKE_CASE__  (					cls						:						List[str]    ,							a						:						Optional[Union[str, os.PathLike]]    ,							**a						:						Dict       ):





         """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(a       ):
             __lowerCamelCase            =    ControlNetModel.from_pretrained(a    ,							**a       )
             controlnets.append(a       )

             idx += 1
             __lowerCamelCase            =    pretrained_model_path + f"""_{idx}"""

         logger.info(f"""{len(a       )} controlnets loaded from {pretrained_model_path}."""       )

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

         return cls(a       )

'''simple docstring'''




from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_tf_available,
    is_tokenizers_available,
    is_torch_available,
)


__UpperCAmelCase					={
    "configuration_mobilebert": [
        "MOBILEBERT_PRETRAINED_CONFIG_ARCHIVE_MAP",
        "MobileBertConfig",
        "MobileBertOnnxConfig",
    ],
    "tokenization_mobilebert": ["MobileBertTokenizer"],
}

try:
  if not is_tokenizers_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  __UpperCAmelCase					=["MobileBertTokenizerFast"]

try:
  if not is_torch_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  __UpperCAmelCase					=[
      "MOBILEBERT_PRETRAINED_MODEL_ARCHIVE_LIST",
      "MobileBertForMaskedLM",
      "MobileBertForMultipleChoice",
      "MobileBertForNextSentencePrediction",
      "MobileBertForPreTraining",
      "MobileBertForQuestionAnswering",
      "MobileBertForSequenceClassification",
      "MobileBertForTokenClassification",
      "MobileBertLayer",
      "MobileBertModel",
      "MobileBertPreTrainedModel",
      "load_tf_weights_in_mobilebert",
  ]

try:
  if not is_tf_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  __UpperCAmelCase					=[
      "TF_MOBILEBERT_PRETRAINED_MODEL_ARCHIVE_LIST",
      "TFMobileBertForMaskedLM",
      "TFMobileBertForMultipleChoice",
      "TFMobileBertForNextSentencePrediction",
      "TFMobileBertForPreTraining",
      "TFMobileBertForQuestionAnswering",
      "TFMobileBertForSequenceClassification",
      "TFMobileBertForTokenClassification",
      "TFMobileBertMainLayer",
      "TFMobileBertModel",
      "TFMobileBertPreTrainedModel",
  ]


if TYPE_CHECKING:
  from .configuration_mobilebert import (
      MOBILEBERT_PRETRAINED_CONFIG_ARCHIVE_MAP,
      MobileBertConfig,
      MobileBertOnnxConfig,
  )
  from .tokenization_mobilebert import MobileBertTokenizer

  try:
    if not is_tokenizers_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .tokenization_mobilebert_fast import MobileBertTokenizerFast

  try:
    if not is_torch_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .modeling_mobilebert import (
        MOBILEBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
        MobileBertForMaskedLM,
        MobileBertForMultipleChoice,
        MobileBertForNextSentencePrediction,
        MobileBertForPreTraining,
        MobileBertForQuestionAnswering,
        MobileBertForSequenceClassification,
        MobileBertForTokenClassification,
        MobileBertLayer,
        MobileBertModel,
        MobileBertPreTrainedModel,
        load_tf_weights_in_mobilebert,
    )

  try:
    if not is_tf_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .modeling_tf_mobilebert import (
        TF_MOBILEBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
        TFMobileBertForMaskedLM,
        TFMobileBertForMultipleChoice,
        TFMobileBertForNextSentencePrediction,
        TFMobileBertForPreTraining,
        TFMobileBertForQuestionAnswering,
        TFMobileBertForSequenceClassification,
        TFMobileBertForTokenClassification,
        TFMobileBertMainLayer,
        TFMobileBertModel,
        TFMobileBertPreTrainedModel,
    )

else:
  import sys

  __UpperCAmelCase					=_LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)
