
"""simple docstring"""




import torch

from diffusers import KDPMaDiscreteScheduler
from diffusers.utils import torch_device

from .test_schedulers import SchedulerCommonTest





class  snake_case     (							__lowerCAmelCase   ):
							UpperCAmelCase__						     =      (KDPMaDiscreteScheduler,)
							UpperCAmelCase__						     =      10

							def _lowercase					(self       ,       **SCREAMING_SNAKE_CASE_    ):

									"""simple docstring"""







									SCREAMING_SNAKE_CASE_						 = {
									    '''num_train_timesteps''': 11_00,
									    '''beta_start''': 0.00_01,
									    '''beta_end''': 0.02,
									    '''beta_schedule''': '''linear''',
									}

									config.update(**lowerCAmelCase_    )
									return config

							def _lowercase					(self    ):

									"""simple docstring"""







									for timesteps in [10, 50, 1_00, 10_00]:
											self.check_over_configs(num_train_timesteps=lowerCAmelCase_    )

							def _lowercase					(self    ):

									"""simple docstring"""







									for beta_start, beta_end in zip([0.0_00_01, 0.00_01, 0.0_01]       ,       [0.00_02, 0.0_02, 0.02]    ):
											self.check_over_configs(beta_start=lowerCAmelCase_       ,       beta_end=lowerCAmelCase_    )

							def _lowercase					(self    ):

									"""simple docstring"""







									for schedule in ["linear", "scaled_linear"]:
											self.check_over_configs(beta_schedule=lowerCAmelCase_    )

							def _lowercase					(self    ):

									"""simple docstring"""







									for prediction_type in ["epsilon", "v_prediction"]:
											self.check_over_configs(prediction_type=lowerCAmelCase_    )

							def _lowercase					(self    ):

									"""simple docstring"""







									SCREAMING_SNAKE_CASE_						 = self.scheduler_classes[0]
									SCREAMING_SNAKE_CASE_						 = self.get_scheduler_config(prediction_type='''v_prediction'''    )
									SCREAMING_SNAKE_CASE_						 = scheduler_class(**lowerCAmelCase_    )

									scheduler.set_timesteps(self.num_inference_steps    )

									SCREAMING_SNAKE_CASE_						 = self.dummy_model()
									SCREAMING_SNAKE_CASE_						 = self.dummy_sample_deter * scheduler.init_noise_sigma
									SCREAMING_SNAKE_CASE_						 = sample.to(lowerCAmelCase_    )

									for i, t in enumerate(scheduler.timesteps    ):
											SCREAMING_SNAKE_CASE_						 = scheduler.scale_model_input(lowerCAmelCase_       ,       lowerCAmelCase_    )

											SCREAMING_SNAKE_CASE_						 = model(lowerCAmelCase_       ,       lowerCAmelCase_    )

											SCREAMING_SNAKE_CASE_						 = scheduler.step(lowerCAmelCase_       ,       lowerCAmelCase_       ,       lowerCAmelCase_    )
											SCREAMING_SNAKE_CASE_						 = output.prev_sample

									SCREAMING_SNAKE_CASE_						 = torch.sum(torch.abs(lowerCAmelCase_    )    )
									SCREAMING_SNAKE_CASE_						 = torch.mean(torch.abs(lowerCAmelCase_    )    )

									if torch_device in ["cpu", "mps"]:
											assert abs(result_sum.item() - 4.6_9_3_4e-0_7    ) < 1e-2
											assert abs(result_mean.item() - 6.1_1_1_2e-1_0    ) < 1e-3
									else:
											# CUDA
											assert abs(result_sum.item() - 4.6_9_3_4_2_8_6_5_0_1_7_0_9_7_2e-0_7    ) < 1e-2
											assert abs(result_mean.item() - 0.00_02    ) < 1e-3

							def _lowercase					(self    ):

									"""simple docstring"""







									if torch_device == "mps":
											return
									SCREAMING_SNAKE_CASE_						 = self.scheduler_classes[0]
									SCREAMING_SNAKE_CASE_						 = self.get_scheduler_config()
									SCREAMING_SNAKE_CASE_						 = scheduler_class(**lowerCAmelCase_    )

									scheduler.set_timesteps(self.num_inference_steps    )

									SCREAMING_SNAKE_CASE_						 = self.dummy_model()
									SCREAMING_SNAKE_CASE_						 = self.dummy_sample_deter * scheduler.init_noise_sigma
									SCREAMING_SNAKE_CASE_						 = sample.to(lowerCAmelCase_    )

									for i, t in enumerate(scheduler.timesteps    ):
											SCREAMING_SNAKE_CASE_						 = scheduler.scale_model_input(lowerCAmelCase_       ,       lowerCAmelCase_    )

											SCREAMING_SNAKE_CASE_						 = model(lowerCAmelCase_       ,       lowerCAmelCase_    )

											SCREAMING_SNAKE_CASE_						 = scheduler.step(lowerCAmelCase_       ,       lowerCAmelCase_       ,       lowerCAmelCase_    )
											SCREAMING_SNAKE_CASE_						 = output.prev_sample

									SCREAMING_SNAKE_CASE_						 = torch.sum(torch.abs(lowerCAmelCase_    )    )
									SCREAMING_SNAKE_CASE_						 = torch.mean(torch.abs(lowerCAmelCase_    )    )

									if torch_device in ["cpu", "mps"]:
											assert abs(result_sum.item() - 20.41_25    ) < 1e-2
											assert abs(result_mean.item() - 0.02_66    ) < 1e-3
									else:
											# CUDA
											assert abs(result_sum.item() - 20.41_25    ) < 1e-2
											assert abs(result_mean.item() - 0.02_66    ) < 1e-3







							def _lowercase					(self    ):

									"""simple docstring"""







									if torch_device == "mps":
											return
									SCREAMING_SNAKE_CASE_						 = self.scheduler_classes[0]
									SCREAMING_SNAKE_CASE_						 = self.get_scheduler_config()
									SCREAMING_SNAKE_CASE_						 = scheduler_class(**lowerCAmelCase_    )

									scheduler.set_timesteps(self.num_inference_steps       ,       device=lowerCAmelCase_    )

									SCREAMING_SNAKE_CASE_						 = self.dummy_model()
									SCREAMING_SNAKE_CASE_						 = self.dummy_sample_deter.to(lowerCAmelCase_    ) * scheduler.init_noise_sigma

									for t in scheduler.timesteps:
											SCREAMING_SNAKE_CASE_						 = scheduler.scale_model_input(lowerCAmelCase_       ,       lowerCAmelCase_    )

											SCREAMING_SNAKE_CASE_						 = model(lowerCAmelCase_       ,       lowerCAmelCase_    )

											SCREAMING_SNAKE_CASE_						 = scheduler.step(lowerCAmelCase_       ,       lowerCAmelCase_       ,       lowerCAmelCase_    )
											SCREAMING_SNAKE_CASE_						 = output.prev_sample

									SCREAMING_SNAKE_CASE_						 = torch.sum(torch.abs(lowerCAmelCase_    )    )
									SCREAMING_SNAKE_CASE_						 = torch.mean(torch.abs(lowerCAmelCase_    )    )

									if str(lowerCAmelCase_    ).startswith('''cpu'''    ):
											# The following sum varies between 148 and 156 on mps. Why?
											assert abs(result_sum.item() - 20.41_25    ) < 1e-2
											assert abs(result_mean.item() - 0.02_66    ) < 1e-3
									else:
											# CUDA
											assert abs(result_sum.item() - 20.41_25    ) < 1e-2
											assert abs(result_mean.item() - 0.02_66    ) < 1e-3

"""simple docstring"""




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        _lowerCamelCase ( __a    ):
		if is_torch_version('''<''',							'''2.0.0'''    ) or not hasattr(__a,							'''_dynamo'''    ):
				return False
		return isinstance(__a,							torch._dynamo.eval_frame.OptimizedModule    )



def        _lowerCamelCase ( __a,							__a = True    ):
		SCREAMING_SNAKE_CASE_						 = (torch.nn.parallel.DistributedDataParallel, torch.nn.DataParallel)

		SCREAMING_SNAKE_CASE_						 = is_compiled_module(__a    )
		if is_compiled:
				SCREAMING_SNAKE_CASE_						 = model
				SCREAMING_SNAKE_CASE_						 = model._orig_mod

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

		while isinstance(__a,							__a    ):
				SCREAMING_SNAKE_CASE_						 = model.module

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

		if is_compiled:
				SCREAMING_SNAKE_CASE_						 = model
				SCREAMING_SNAKE_CASE_						 = compiled_model

		return model



def        _lowerCamelCase ( ):
		PartialState().wait_for_everyone()



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



@contextmanager
def        _lowerCamelCase ( **__a    ):
		for key, value in kwargs.items():
				SCREAMING_SNAKE_CASE_						 = str(__a    )

		yield

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



def        _lowerCamelCase ( __a    ):
		if not hasattr(__a,							'''__qualname__'''    ) and not hasattr(__a,							'''__name__'''    ):
				SCREAMING_SNAKE_CASE_						 = getattr(__a,							'''__class__''',							__a    )
		if hasattr(__a,							'''__qualname__'''    ):
				return obj.__qualname__
		if hasattr(__a,							'''__name__'''    ):
				return obj.__name__
		return str(__a    )



def        _lowerCamelCase ( __a,							__a    ):
		for key, value in source.items():
				if isinstance(__a,							__a    ):
						SCREAMING_SNAKE_CASE_						 = destination.setdefault(__a,							{}    )
						merge_dicts(__a,							__a    )
				else:
						SCREAMING_SNAKE_CASE_						 = value

		return destination



def        _lowerCamelCase ( __a = None    ):
		if port is None:
				SCREAMING_SNAKE_CASE_						 = 29_500
		with socket.socket(socket.AF_INET,							socket.SOCK_STREAM    ) as s:
				return s.connect_ex(('''localhost''', port)    ) == 0