



'''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__       :	Any    = logging.get_logger(__name__)






class 						_UpperCamelCase  (							lowerCamelCase__    ):


			'''simple docstring'''

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





					"""simple docstring"""






					super().__init__()
					__SCREAMING_SNAKE_CASE       :		List[str]           =		nn.ModuleList(lowerCAmelCase__						)

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





					"""simple docstring"""






					for i, (image, scale, controlnet) in enumerate(zip(lowerCAmelCase__							,     lowerCAmelCase__							,     self.nets						)						):
							__SCREAMING_SNAKE_CASE       ,						__SCREAMING_SNAKE_CASE       :		Optional[Any]           =		controlnet(
							    lowerCAmelCase__							,     lowerCAmelCase__							,     lowerCAmelCase__							,     lowerCAmelCase__							,     lowerCAmelCase__							,     lowerCAmelCase__							,     lowerCAmelCase__							,     lowerCAmelCase__							,     lowerCAmelCase__							,     lowerCAmelCase__							,     lowerCAmelCase__							,     )

							# merge samples
							if i == 0:
									__SCREAMING_SNAKE_CASE       ,						__SCREAMING_SNAKE_CASE       :		Union[str, Any]           =		down_samples, mid_sample
							else:
									__SCREAMING_SNAKE_CASE       :		str           =		[
									    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 UpperCamelCase__					(							self     : Union[str, Any]							,     lowerCAmelCase__     : Union[str, os.PathLike]							,     lowerCAmelCase__     : bool = True							,     lowerCAmelCase__     : Callable = None							,     lowerCAmelCase__     : bool = False							,     lowerCAmelCase__     : Optional[str] = None							,     ):





					"""simple docstring"""






					__SCREAMING_SNAKE_CASE       :		Union[str, Any]           =		0
					__SCREAMING_SNAKE_CASE       :		Optional[int]           =		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       :		List[str]           =		model_path_to_save + F"_{idx}"



			@classmethod
			def UpperCamelCase__					(							cls     : str							,     lowerCAmelCase__     : Optional[Union[str, os.PathLike]]							,     **lowerCAmelCase__     : List[str]						):





					"""simple docstring"""






					__SCREAMING_SNAKE_CASE       :		str           =		0
					__SCREAMING_SNAKE_CASE       :		Union[str, Any]           =		[]

					# 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       :		Tuple           =		pretrained_model_path
					while os.path.isdir(lowerCAmelCase__						):
							__SCREAMING_SNAKE_CASE       :		Union[str, Any]           =		ControlNetModel.from_pretrained(lowerCAmelCase__							,     **lowerCAmelCase__						)
							controlnets.append(lowerCAmelCase__						)

							idx += 1
							__SCREAMING_SNAKE_CASE       :		Any           =		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'''




# Copyright 2021 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import argparse

from ...utils.dataclasses import (
    ComputeEnvironment,
    DistributedType,
    DynamoBackend,
    PrecisionType,
    SageMakerDistributedType,
)
from ..menu import BulletMenu


UpperCamelCase__       :	List[Any]    = [
    '''EAGER''',
    '''AOT_EAGER''',
    '''INDUCTOR''',
    '''NVFUSER''',
    '''AOT_NVFUSER''',
    '''AOT_CUDAGRAPHS''',
    '''OFI''',
    '''FX2TRT''',
    '''ONNXRT''',
    '''IPEX''',
]


def 				lowerCAmelCase_   (	_lowerCamelCase:   str       ,      _lowerCamelCase:   Union[str, Any]=None       ,      _lowerCamelCase:   Optional[int]=None       ,      _lowerCamelCase:   str=None					):
		__SCREAMING_SNAKE_CASE       :		Optional[int]           =		True
		while ask_again:
				__SCREAMING_SNAKE_CASE       :		Tuple           =		input(_lowerCamelCase					)
				try:
						if default is not None and len(_lowerCamelCase					) == 0:
								return default
						return convert_value(_lowerCamelCase					) if convert_value is not None else result
				except Exception:
						if error_message is not None:
								print(_lowerCamelCase					)


def 				lowerCAmelCase_   (	_lowerCamelCase:   Tuple       ,      _lowerCamelCase:   Union[str, Any]=[]       ,      _lowerCamelCase:   List[Any]=None       ,      _lowerCamelCase:   Optional[Any]=0					):
		__SCREAMING_SNAKE_CASE       :		Union[str, Any]           =		BulletMenu(_lowerCamelCase       ,      _lowerCamelCase					)
		__SCREAMING_SNAKE_CASE       :		Dict           =		menu.run(default_choice=_lowerCamelCase					)
		return convert_value(_lowerCamelCase					) if convert_value is not None else result


def 				lowerCAmelCase_   (	_lowerCamelCase:   Optional[Any]					):
		__SCREAMING_SNAKE_CASE       :		List[str]           =		int(_lowerCamelCase					)
		return ComputeEnvironment(["""LOCAL_MACHINE""", """AMAZON_SAGEMAKER"""][value]					)


def 				lowerCAmelCase_   (	_lowerCamelCase:   Any					):
		__SCREAMING_SNAKE_CASE       :		str           =		int(_lowerCamelCase					)
		return DistributedType(["""NO""", """MULTI_CPU""", """MULTI_XPU""", """MULTI_GPU""", """MULTI_NPU""", """TPU"""][value]					)


def 				lowerCAmelCase_   (	_lowerCamelCase:   Tuple					):
		__SCREAMING_SNAKE_CASE       :		Tuple           =		int(_lowerCamelCase					)
		return DynamoBackend(DYNAMO_BACKENDS[value]					).value


def 				lowerCAmelCase_   (	_lowerCamelCase:   Union[str, Any]					):
		__SCREAMING_SNAKE_CASE       :		List[str]           =		int(_lowerCamelCase					)
		return PrecisionType(["""no""", """fp16""", """bf16""", """fp8"""][value]					)


def 				lowerCAmelCase_   (	_lowerCamelCase:   Tuple					):
		__SCREAMING_SNAKE_CASE       :		int           =		int(_lowerCamelCase					)
		return SageMakerDistributedType(["""NO""", """DATA_PARALLEL""", """MODEL_PARALLEL"""][value]					)


def 				lowerCAmelCase_   (	_lowerCamelCase:   List[Any]					):
		return {"yes": True, "no": False}[value.lower()]






class 						_UpperCamelCase  (							argparse.RawDescriptionHelpFormatter    ):


			'''simple docstring'''

			def UpperCamelCase__					(							self     : Tuple							,     lowerCAmelCase__     : Optional[Any]							,     lowerCAmelCase__     : Dict							,     lowerCAmelCase__     : List[Any]							,     lowerCAmelCase__     : List[Any]						):





					"""simple docstring"""






					__SCREAMING_SNAKE_CASE       :		Union[str, Any]           =		super()._format_usage(lowerCAmelCase__							,     lowerCAmelCase__							,     lowerCAmelCase__							,     lowerCAmelCase__						)
					__SCREAMING_SNAKE_CASE       :		Any           =		usage.replace("""<command> [<args>] """							,     """"""						)
					return usage