




import importlib.util
import json
import os
import warnings
from dataclasses import dataclass, field

import torch

from ..training_args import TrainingArguments
from ..utils import cached_property, is_sagemaker_dp_enabled, logging


SCREAMING_SNAKE_CASE__     :     List[Any]		    =       logging.get_logger(__name__)

def lowercase    (  )			->							Any:


			'''simple docstring'''







			SCREAMING_SNAKE_CASE_								=						os.getenv('SM_HP_MP_PARAMETERS'       ,    '{}' )
			try:
						# Parse it and check the field "partitions" is included, it is required for model parallel.
						SCREAMING_SNAKE_CASE_								=						json.loads(SCREAMING_SNAKE_CASE )
						if "partitions" not in smp_options:
									return False
			except json.JSONDecodeError:
						return False

			# Get the sagemaker specific framework parameters from mpi_options variable.
			SCREAMING_SNAKE_CASE_								=						os.getenv('SM_FRAMEWORK_PARAMS'       ,    '{}' )
			try:
						# Parse it and check the field "sagemaker_distributed_dataparallel_enabled".
						SCREAMING_SNAKE_CASE_								=						json.loads(SCREAMING_SNAKE_CASE )
						if not mpi_options.get('sagemaker_mpi_enabled'       ,    SCREAMING_SNAKE_CASE ):
									return False
			except json.JSONDecodeError:
						return False
			# Lastly, check if the `smdistributed` module is present.
			return importlib.util.find_spec('smdistributed' ) is not None


if is_sagemaker_model_parallel_available():
				import smdistributed.modelparallel.torch as smp

				smp.init()



@dataclass
class       a_       (     SCREAMING_SNAKE_CASE__     ):
	A    							= field(
	    default=''''''  ,    metadata={'''help''': '''Used by the SageMaker launcher to send mp-specific args. Ignored in SageMakerTrainer'''}  ,    )





	def 						A_(							self      )  ->					Optional[int]:




				"""simple docstring"""



				super().__post_init__()
				warnings.warn(
				    '`SageMakerTrainingArguments` is deprecated and will be removed in v5 of Transformers. You can use '
				    '`TrainingArguments` instead.'      ,     SCREAMING_SNAKE_CASE      ,     )





	@cached_property
	def 						A_(							self      )  ->					"torch.device":




				"""simple docstring"""



				logger.info('PyTorch: setting up devices'      )
				if torch.distributed.is_available() and torch.distributed.is_initialized() and self.local_rank == -1:
							logger.warning(
							    'torch.distributed process group is initialized, but local_rank == -1. '
							    'In order to use Torch DDP, launch your script with `python -m torch.distributed.launch'      )
				if self.no_cuda:
							SCREAMING_SNAKE_CASE_								=						torch.device('cpu'      )
							SCREAMING_SNAKE_CASE_								=						0
				elif is_sagemaker_model_parallel_available():
							SCREAMING_SNAKE_CASE_								=						smp.local_rank()
							SCREAMING_SNAKE_CASE_								=						torch.device('cuda'      ,     SCREAMING_SNAKE_CASE      )
							SCREAMING_SNAKE_CASE_								=						1
				elif is_sagemaker_dp_enabled():
							import smdistributed.dataparallel.torch.torch_smddp  # noqa: F401

							torch.distributed.init_process_group(backend='smddp'      ,     timeout=self.ddp_timeout_delta      )
							SCREAMING_SNAKE_CASE_								=						int(os.getenv('SMDATAPARALLEL_LOCAL_RANK'      )      )
							SCREAMING_SNAKE_CASE_								=						torch.device('cuda'      ,     self.local_rank      )
							SCREAMING_SNAKE_CASE_								=						1
				elif self.local_rank == -1:
							# if n_gpu is > 1 we'll use nn.DataParallel.
							# If you only want to use a specific subset of GPUs use `CUDA_VISIBLE_DEVICES=0`
							# Explicitly set CUDA to the first (index 0) CUDA device, otherwise `set_device` will
							# trigger an error that a device index is missing. Index 0 takes into account the
							# GPUs available in the environment, so `CUDA_VISIBLE_DEVICES=1,2` with `cuda:0`
							# will use the first GPU in that env, i.e. GPU#1
							SCREAMING_SNAKE_CASE_								=						torch.device('cuda:0' if torch.cuda.is_available() else 'cpu'      )
							# Sometimes the line in the postinit has not been run before we end up here, so just checking we're not at
							# the default value.
							SCREAMING_SNAKE_CASE_								=						torch.cuda.device_count()
				else:
							# Here, we'll use torch.distributed.
							# Initializes the distributed backend which will take care of synchronizing nodes/GPUs
							if not torch.distributed.is_initialized():
										torch.distributed.init_process_group(backend='nccl'      ,     timeout=self.ddp_timeout_delta      )
							SCREAMING_SNAKE_CASE_								=						torch.device('cuda'      ,     self.local_rank      )
							SCREAMING_SNAKE_CASE_								=						1

				if device.type == "cuda":
							torch.cuda.set_device(SCREAMING_SNAKE_CASE      )

				return device





	@property
	def 						A_(							self      )  ->					Union[str, Any]:




				"""simple docstring"""



				if is_sagemaker_model_parallel_available():
							return smp.dp_size()

				return super().world_size





	@property
	def 						A_(							self      )  ->					int:




				"""simple docstring"""



				return not is_sagemaker_model_parallel_available()





	@property
	def 						A_(							self      )  ->					Optional[int]:




				"""simple docstring"""



				return False






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__     :     Union[str, Any]		    =       logging.get_logger(__name__)



class       a_       (     SCREAMING_SNAKE_CASE__     ):





	def __init__(							self      ,     SCREAMING_SNAKE_CASE      )  ->					Union[str, Any]:




				"""simple docstring"""



				super().__init__()
				SCREAMING_SNAKE_CASE_								=						nn.ModuleList(SCREAMING_SNAKE_CASE      )





	def 						A_(							self      ,     SCREAMING_SNAKE_CASE      ,     SCREAMING_SNAKE_CASE      ,     SCREAMING_SNAKE_CASE      ,     SCREAMING_SNAKE_CASE      ,     SCREAMING_SNAKE_CASE      ,     SCREAMING_SNAKE_CASE = None      ,     SCREAMING_SNAKE_CASE = None      ,     SCREAMING_SNAKE_CASE = None      ,     SCREAMING_SNAKE_CASE = None      ,     SCREAMING_SNAKE_CASE = False      ,     SCREAMING_SNAKE_CASE = True      ,     )  ->					Union[ControlNetOutput, Tuple]:




				"""simple docstring"""



				for i, (image, scale, controlnet) in enumerate(zip(SCREAMING_SNAKE_CASE      ,     SCREAMING_SNAKE_CASE      ,     self.nets      )      ):
							SCREAMING_SNAKE_CASE_						,		SCREAMING_SNAKE_CASE_								=						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:
										SCREAMING_SNAKE_CASE_						,		SCREAMING_SNAKE_CASE_								=						down_samples, mid_sample
							else:
										SCREAMING_SNAKE_CASE_								=						[
										    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 						A_(							self      ,     SCREAMING_SNAKE_CASE      ,     SCREAMING_SNAKE_CASE = True      ,     SCREAMING_SNAKE_CASE = None      ,     SCREAMING_SNAKE_CASE = False      ,     SCREAMING_SNAKE_CASE = None      ,     )  ->					int:




				"""simple docstring"""



				SCREAMING_SNAKE_CASE_								=						0
				SCREAMING_SNAKE_CASE_								=						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
							SCREAMING_SNAKE_CASE_								=						model_path_to_save + f'_{idx}'





	@classmethod
	def 						A_(							cls      ,     SCREAMING_SNAKE_CASE      ,     **SCREAMING_SNAKE_CASE      )  ->					Tuple:




				"""simple docstring"""



				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(SCREAMING_SNAKE_CASE      ):
							SCREAMING_SNAKE_CASE_								=						ControlNetModel.from_pretrained(SCREAMING_SNAKE_CASE      ,     **SCREAMING_SNAKE_CASE      )
							controlnets.append(SCREAMING_SNAKE_CASE      )

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