


import os
import socket
from contextlib import contextmanager

import torch

from ..commands.config.default import write_basic_config  # noqa: F401
from ..state import PartialState
from .dataclasses import DistributedType
from .imports import is_deepspeed_available, is_tpu_available
from .transformer_engine import convert_model
from .versions import is_torch_version


if is_deepspeed_available():
							from deepspeed import DeepSpeedEngine

if is_tpu_available(check_device=False):
							import torch_xla.core.xla_model as xm


def 			UpperCAmelCase_				(      UpperCAmelCase__  ):
					if is_torch_version("""<"""						,							"""2.0.0"""  ) or not hasattr(UpperCAmelCase__						,							"""_dynamo"""  ):
										return False
					return isinstance(UpperCAmelCase__						,							torch._dynamo.eval_frame.OptimizedModule  )


def 			UpperCAmelCase_				(      UpperCAmelCase__						,							UpperCAmelCase__ = True  ):
					lowercase_  							=			(torch.nn.parallel.DistributedDataParallel, torch.nn.DataParallel)

					lowercase_  							=			is_compiled_module(UpperCAmelCase__  )
					if is_compiled:
										lowercase_  							=			model
										lowercase_  							=			model._orig_mod

					if is_deepspeed_available():
										options += (DeepSpeedEngine,)

					while isinstance(UpperCAmelCase__						,							UpperCAmelCase__  ):
										lowercase_  							=			model.module

					if not keep_fpaa_wrapper:
										lowercase_  							=			getattr(UpperCAmelCase__						,							"""forward"""  )
										lowercase_  							=			model.__dict__.pop("""_original_forward"""						,							UpperCAmelCase__  )
										if original_forward is not None:
															while hasattr(UpperCAmelCase__						,							"""__wrapped__"""  ):
																				lowercase_  							=			forward.__wrapped__
																				if forward == original_forward:
																									break
															lowercase_  							=			forward
										if getattr(UpperCAmelCase__						,							"""_converted_to_transformer_engine"""						,							UpperCAmelCase__  ):
															convert_model(UpperCAmelCase__						,							to_transformer_engine=UpperCAmelCase__  )

					if is_compiled:
										lowercase_  							=			model
										lowercase_  							=			compiled_model

					return model


def 			UpperCAmelCase_				(      ):
					PartialState().wait_for_everyone()


def 			UpperCAmelCase_				(      UpperCAmelCase__						,							UpperCAmelCase__  ):
					if PartialState().distributed_type == DistributedType.TPU:
										xm.save(UpperCAmelCase__						,							UpperCAmelCase__  )
					elif PartialState().local_process_index == 0:
										torch.save(UpperCAmelCase__						,							UpperCAmelCase__  )


@contextmanager
def 			UpperCAmelCase_				(      **UpperCAmelCase__  ):
					for key, value in kwargs.items():
										lowercase_  							=			str(UpperCAmelCase__  )

					yield

					for key in kwargs:
										if key.upper() in os.environ:
															del os.environ[key.upper()]


def 			UpperCAmelCase_				(      UpperCAmelCase__  ):
					if not hasattr(UpperCAmelCase__						,							"""__qualname__"""  ) and not hasattr(UpperCAmelCase__						,							"""__name__"""  ):
										lowercase_  							=			getattr(UpperCAmelCase__						,							"""__class__"""						,							UpperCAmelCase__  )
					if hasattr(UpperCAmelCase__						,							"""__qualname__"""  ):
										return obj.__qualname__
					if hasattr(UpperCAmelCase__						,							"""__name__"""  ):
										return obj.__name__
					return str(UpperCAmelCase__  )


def 			UpperCAmelCase_				(      UpperCAmelCase__						,							UpperCAmelCase__  ):
					for key, value in source.items():
										if isinstance(UpperCAmelCase__						,							UpperCAmelCase__  ):
															lowercase_  							=			destination.setdefault(UpperCAmelCase__						,							{}  )
															merge_dicts(UpperCAmelCase__						,							UpperCAmelCase__  )
										else:
															lowercase_  							=			value

					return destination



def 			UpperCAmelCase_				(      UpperCAmelCase__ = None  ):
					if port is None:
										lowercase_  							=			2_9_5_0_0
					with socket.socket(socket.AF_INET						,							socket.SOCK_STREAM  ) as s:
										return s.connect_ex(("""localhost""", port)  ) == 0






from ...utils import (
    OptionalDependencyNotAvailable,
    is_flax_available,
    is_torch_available,
    is_transformers_available,
)


try:
							if not (is_transformers_available() and is_torch_available()):
														raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
							from ...utils.dummy_torch_and_transformers_objects import *  # noqa F403
else:
							from .multicontrolnet import MultiControlNetModel
							from .pipeline_controlnet import StableDiffusionControlNetPipeline
							from .pipeline_controlnet_imgaimg import StableDiffusionControlNetImgaImgPipeline
							from .pipeline_controlnet_inpaint import StableDiffusionControlNetInpaintPipeline


if is_transformers_available() and is_flax_available():
							from .pipeline_flax_controlnet import FlaxStableDiffusionControlNetPipeline


