






"""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   :							List[Any]          =		logging.get_logger(__name__)






class 	a		(   __snake_case						):






    def __init__(			self	:				Union[str, Any]			,		__SCREAMING_SNAKE_CASE	:				Union[List[ControlNetModel], Tuple[ControlNetModel]]   )	->						List[Any]:
      super().__init__()
      lowerCamelCase_       	=		nn.ModuleList(__SCREAMING_SNAKE_CASE   )






    def   UpperCamelCase			(			self	:				Union[str, Any]			,		__SCREAMING_SNAKE_CASE	:				torch.FloatTensor			,		__SCREAMING_SNAKE_CASE	:				Union[torch.Tensor, float, int]			,		__SCREAMING_SNAKE_CASE	:				torch.Tensor			,		__SCREAMING_SNAKE_CASE	:				List[torch.tensor]			,		__SCREAMING_SNAKE_CASE	:				List[float]			,		__SCREAMING_SNAKE_CASE	:				Optional[torch.Tensor] = None			,		__SCREAMING_SNAKE_CASE	:				Optional[torch.Tensor] = None			,		__SCREAMING_SNAKE_CASE	:				Optional[torch.Tensor] = None			,		__SCREAMING_SNAKE_CASE	:				Optional[Dict[str, Any]] = None			,		__SCREAMING_SNAKE_CASE	:				bool = False			,		__SCREAMING_SNAKE_CASE	:				bool = True			,		)	->						Union[ControlNetOutput, Tuple]:
      for i, (image, scale, controlnet) in enumerate(zip(__SCREAMING_SNAKE_CASE			,		__SCREAMING_SNAKE_CASE			,		self.nets   )   ):
        lowerCamelCase_	,      lowerCamelCase_       	=		controlnet(
            __SCREAMING_SNAKE_CASE			,		__SCREAMING_SNAKE_CASE			,		__SCREAMING_SNAKE_CASE			,		__SCREAMING_SNAKE_CASE			,		__SCREAMING_SNAKE_CASE			,		__SCREAMING_SNAKE_CASE			,		__SCREAMING_SNAKE_CASE			,		__SCREAMING_SNAKE_CASE			,		__SCREAMING_SNAKE_CASE			,		__SCREAMING_SNAKE_CASE			,		__SCREAMING_SNAKE_CASE			,		)

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

      return down_block_res_samples, mid_block_res_sample






    def   UpperCamelCase			(			self	:				Tuple			,		__SCREAMING_SNAKE_CASE	:				Union[str, os.PathLike]			,		__SCREAMING_SNAKE_CASE	:				bool = True			,		__SCREAMING_SNAKE_CASE	:				Callable = None			,		__SCREAMING_SNAKE_CASE	:				bool = False			,		__SCREAMING_SNAKE_CASE	:				Optional[str] = None			,		)	->						int:
      lowerCamelCase_       	=		0
      lowerCamelCase_       	=		save_directory
      for controlnet in self.nets:
        controlnet.save_pretrained(
            __SCREAMING_SNAKE_CASE			,		is_main_process=__SCREAMING_SNAKE_CASE			,		save_function=__SCREAMING_SNAKE_CASE			,		safe_serialization=__SCREAMING_SNAKE_CASE			,		variant=__SCREAMING_SNAKE_CASE			,		)

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






    @classmethod
    def   UpperCamelCase			(			cls	:				Union[str, Any]			,		__SCREAMING_SNAKE_CASE	:				Optional[Union[str, os.PathLike]]			,		**__SCREAMING_SNAKE_CASE	:				Optional[Any]   )	->						Optional[Any]:
      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(__SCREAMING_SNAKE_CASE   ):
        lowerCamelCase_       	=		ControlNetModel.from_pretrained(__SCREAMING_SNAKE_CASE			,		**__SCREAMING_SNAKE_CASE   )
        controlnets.append(__SCREAMING_SNAKE_CASE   )

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

      logger.info(F'''{len(__SCREAMING_SNAKE_CASE   )} controlnets loaded from {pretrained_model_path}.'''   )

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

      return cls(__SCREAMING_SNAKE_CASE   )












"""simple docstring"""





from typing import TYPE_CHECKING

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


_SCREAMING_SNAKE_CASE   :							List[str]          =		{}

try:
  if not is_sentencepiece_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  _SCREAMING_SNAKE_CASE   :							str          =		['''NllbTokenizer''']

try:
  if not is_tokenizers_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  _SCREAMING_SNAKE_CASE   :							Optional[Any]          =		['''NllbTokenizerFast''']


if TYPE_CHECKING:
  try:
    if not is_sentencepiece_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .tokenization_nllb import NllbTokenizer

  try:
    if not is_tokenizers_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .tokenization_nllb_fast import NllbTokenizerFast

else:
  import sys

  _SCREAMING_SNAKE_CASE   :							Any          =		_LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)




