


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


_lowerCamelCase     =  logging.get_logger(__name__)


class 			lowerCamelCase_		(			_UpperCAmelCase   ):





				"""simple docstring"""







				def __init__(						self       ,		UpperCAmelCase__     ):
								super().__init__()
								SCREAMING_SNAKE_CASE__     							=				nn.ModuleList(lowerCamelCase_     )



				def   lowerCAmelCase__		(						self       ,		UpperCAmelCase__       ,		UpperCAmelCase__       ,		UpperCAmelCase__       ,		UpperCAmelCase__       ,		UpperCAmelCase__       ,		UpperCAmelCase__ = None       ,		UpperCAmelCase__ = None       ,		UpperCAmelCase__ = None       ,		UpperCAmelCase__ = None       ,		UpperCAmelCase__ = False       ,		UpperCAmelCase__ = True       ,		):




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

												# merge samples
												if i == 0:
																SCREAMING_SNAKE_CASE__     							=				down_samples, mid_sample
												else:
																SCREAMING_SNAKE_CASE__     							=				[
																    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   lowerCAmelCase__		(						self       ,		UpperCAmelCase__       ,		UpperCAmelCase__ = True       ,		UpperCAmelCase__ = None       ,		UpperCAmelCase__ = False       ,		UpperCAmelCase__ = None       ,		):
								SCREAMING_SNAKE_CASE__     							=				0
								SCREAMING_SNAKE_CASE__     							=				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
												SCREAMING_SNAKE_CASE__     							=				model_path_to_save + f'''_{idx}'''






				@classmethod
				def   lowerCAmelCase__		(						cls       ,		UpperCAmelCase__       ,		**UpperCAmelCase__     ):
								SCREAMING_SNAKE_CASE__     							=				0
								SCREAMING_SNAKE_CASE__     							=				[]

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

												idx += 1
												SCREAMING_SNAKE_CASE__     							=				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_     )




"""simple docstring"""






from typing import TYPE_CHECKING

# rely on isort to merge the imports
from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available


_lowerCamelCase     =  {
    'configuration_autoformer': [
        'AUTOFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP',
        'AutoformerConfig',
    ],
}

try:
	if not is_torch_available():
		raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
	pass
else:
	_lowerCamelCase     =  [
	    'AUTOFORMER_PRETRAINED_MODEL_ARCHIVE_LIST',
	    'AutoformerForPrediction',
	    'AutoformerModel',
	    'AutoformerPreTrainedModel',
	]


if TYPE_CHECKING:
	from .configuration_autoformer import (
	    AUTOFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP,
	    AutoformerConfig,
	)

	try:
		if not is_torch_available():
			raise OptionalDependencyNotAvailable()
	except OptionalDependencyNotAvailable:
		pass
	else:
		from .modeling_autoformer import (
		    AUTOFORMER_PRETRAINED_MODEL_ARCHIVE_LIST,
		    AutoformerForPrediction,
		    AutoformerModel,
		    AutoformerPreTrainedModel,
		)

else:
	import sys

	_lowerCamelCase     =  _LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)
