






import gc
import random
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import AutoencoderKL, CycleDiffusionPipeline, DDIMScheduler, UNetaDConditionModel
from diffusers.utils import floats_tensor, load_image, load_numpy, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu, skip_mps

from ..pipeline_params import (
    IMAGE_TO_IMAGE_IMAGE_PARAMS,
    TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS,
    TEXT_GUIDED_IMAGE_VARIATION_PARAMS,
)
from ..test_pipelines_common import PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()

class     _SCREAMING_SNAKE_CASE       (			_a		,   _a		,   unittest.TestCase      ):
	snake_case__							:     Optional[int]     	=  CycleDiffusionPipeline
	snake_case__							:     Union[str, Any]     	=  TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {
	    """negative_prompt""",
	    """height""",
	    """width""",
	    """negative_prompt_embeds""",
	}
	snake_case__							:     Any     	=  PipelineTesterMixin.required_optional_params - {"""latents"""}
	snake_case__							:     Tuple     	=  TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS.union({"""source_prompt"""}      )
	snake_case__							:     Tuple     	=  IMAGE_TO_IMAGE_IMAGE_PARAMS
	snake_case__							:     List[Any]     	=  IMAGE_TO_IMAGE_IMAGE_PARAMS





	def 						_A (		self				:					str	):
		torch.manual_seed(0	)
		UpperCamelCase     :Optional[int]          =	UNetaDConditionModel(
		    block_out_channels=(32, 64) ,  layers_per_block=2 ,  sample_size=32 ,  in_channels=4 ,  out_channels=4 ,  down_block_types=("""DownBlock2D""", """CrossAttnDownBlock2D""") ,  up_block_types=("""CrossAttnUpBlock2D""", """UpBlock2D""") ,  cross_attention_dim=32 ,  )
		UpperCamelCase     :Union[str, Any]          =	DDIMScheduler(
		    beta_start=0.00085 ,  beta_end=0.012 ,  beta_schedule="""scaled_linear""" ,  num_train_timesteps=1_000 ,  clip_sample=__lowerCamelCase ,  set_alpha_to_one=__lowerCamelCase ,  )
		torch.manual_seed(0	)
		UpperCamelCase     :Dict          =	AutoencoderKL(
		    block_out_channels=[32, 64] ,  in_channels=3 ,  out_channels=3 ,  down_block_types=["""DownEncoderBlock2D""", """DownEncoderBlock2D"""] ,  up_block_types=["""UpDecoderBlock2D""", """UpDecoderBlock2D"""] ,  latent_channels=4 ,  )
		torch.manual_seed(0	)
		UpperCamelCase     :List[Any]          =	CLIPTextConfig(
		    bos_token_id=0 ,  eos_token_id=2 ,  hidden_size=32 ,  intermediate_size=37 ,  layer_norm_eps=1E-05 ,  num_attention_heads=4 ,  num_hidden_layers=5 ,  pad_token_id=1 ,  vocab_size=1_000 ,  )
		UpperCamelCase     :List[Any]          =	CLIPTextModel(__lowerCamelCase	)
		UpperCamelCase     :Optional[int]          =	CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""	)

		UpperCamelCase     :List[str]          =	{
		    """unet""": unet,
		    """scheduler""": scheduler,
		    """vae""": vae,
		    """text_encoder""": text_encoder,
		    """tokenizer""": tokenizer,
		    """safety_checker""": None,
		    """feature_extractor""": None,
		}
		return components





	def 						_A (		self				:					List[Any] ,  __lowerCamelCase				:					int ,  __lowerCamelCase				:					Union[str, Any]=0	):
		UpperCamelCase     :Union[str, Any]          =	floats_tensor((1, 3, 32, 32) ,  rng=random.Random(__lowerCamelCase	)	).to(__lowerCamelCase	)
		UpperCamelCase     :List[str]          =	image / 2 + 0.5
		if str(__lowerCamelCase	).startswith("""mps"""	):
			UpperCamelCase     :Union[str, Any]          =	torch.manual_seed(__lowerCamelCase	)
		else:
			UpperCamelCase     :List[Any]          =	torch.Generator(device=__lowerCamelCase	).manual_seed(__lowerCamelCase	)
		UpperCamelCase     :Optional[Any]          =	{
		    """prompt""": """An astronaut riding an elephant""",
		    """source_prompt""": """An astronaut riding a horse""",
		    """image""": image,
		    """generator""": generator,
		    """num_inference_steps""": 2,
		    """eta""": 0.1,
		    """strength""": 0.8,
		    """guidance_scale""": 3,
		    """source_guidance_scale""": 1,
		    """output_type""": """numpy""",
		}
		return inputs





	def 						_A (		self				:					Optional[int]	):
		UpperCamelCase     :Dict          =	"""cpu"""  # ensure determinism for the device-dependent torch.Generator

		UpperCamelCase     :str          =	self.get_dummy_components()
		UpperCamelCase     :int          =	CycleDiffusionPipeline(**__lowerCamelCase	)
		UpperCamelCase     :Dict          =	pipe.to(__lowerCamelCase	)
		pipe.set_progress_bar_config(disable=__lowerCamelCase	)

		UpperCamelCase     :Tuple          =	self.get_dummy_inputs(__lowerCamelCase	)
		UpperCamelCase     :Any          =	pipe(**__lowerCamelCase	)
		UpperCamelCase     :Dict          =	output.images

		UpperCamelCase     :Any          =	images[0, -3:, -3:, -1]

		assert images.shape == (1, 32, 32, 3)
		UpperCamelCase     :int          =	np.array([0.4459, 0.4943, 0.4544, 0.6643, 0.5474, 0.4327, 0.5701, 0.5959, 0.5179]	)

		assert np.abs(image_slice.flatten() - expected_slice	).max() < 1E-2





	@unittest.skipIf(torch_device != """cuda""" ,  """This test requires a GPU"""	)
	def 						_A (		self				:					Optional[int]	):
		UpperCamelCase     :Dict          =	self.get_dummy_components()
		for name, module in components.items():
			if hasattr(__lowerCamelCase ,  """half"""	):
				UpperCamelCase     :Optional[Any]          =	module.half()
		UpperCamelCase     :int          =	CycleDiffusionPipeline(**__lowerCamelCase	)
		UpperCamelCase     :Optional[Any]          =	pipe.to(__lowerCamelCase	)
		pipe.set_progress_bar_config(disable=__lowerCamelCase	)

		UpperCamelCase     :Any          =	self.get_dummy_inputs(__lowerCamelCase	)
		UpperCamelCase     :Optional[int]          =	pipe(**__lowerCamelCase	)
		UpperCamelCase     :List[str]          =	output.images

		UpperCamelCase     :Optional[Any]          =	images[0, -3:, -3:, -1]

		assert images.shape == (1, 32, 32, 3)
		UpperCamelCase     :List[str]          =	np.array([0.3506, 0.4543, 0.446, 0.4575, 0.5195, 0.4155, 0.5273, 0.518, 0.4116]	)

		assert np.abs(image_slice.flatten() - expected_slice	).max() < 1E-2





	@skip_mps
	def 						_A (		self				:					List[str]	):
		return super().test_save_load_local()





	@unittest.skip("""non-deterministic pipeline"""	)
	def 						_A (		self				:					Optional[int]	):
		return super().test_inference_batch_single_identical()





	@skip_mps
	def 						_A (		self				:					Optional[int]	):
		return super().test_dict_tuple_outputs_equivalent()





	@skip_mps
	def 						_A (		self				:					List[str]	):
		return super().test_save_load_optional_components()





	@skip_mps
	def 						_A (		self				:					Tuple	):
		return super().test_attention_slicing_forward_pass()

@slow
@require_torch_gpu
class     _SCREAMING_SNAKE_CASE       (			unittest.TestCase      ):





	def 						_A (		self				:					Optional[int]	):
		# clean up the VRAM after each test
		super().tearDown()
		gc.collect()
		torch.cuda.empty_cache()





	def 						_A (		self				:					Tuple	):
		UpperCamelCase     :int          =	load_image(
		    """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"""
		    """/cycle-diffusion/black_colored_car.png"""	)
		UpperCamelCase     :List[Any]          =	load_numpy(
		    """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/cycle-diffusion/blue_colored_car_fp16.npy"""	)
		UpperCamelCase     :Any          =	init_image.resize((512, 512)	)

		UpperCamelCase     :Dict          =	"""CompVis/stable-diffusion-v1-4"""
		UpperCamelCase     :List[Any]          =	DDIMScheduler.from_pretrained(__lowerCamelCase ,  subfolder="""scheduler"""	)
		UpperCamelCase     :str          =	CycleDiffusionPipeline.from_pretrained(
		    __lowerCamelCase ,  scheduler=__lowerCamelCase ,  safety_checker=__lowerCamelCase ,  torch_dtype=torch.floataa ,  revision="""fp16"""	)

		pipe.to(__lowerCamelCase	)
		pipe.set_progress_bar_config(disable=__lowerCamelCase	)
		pipe.enable_attention_slicing()

		UpperCamelCase     :List[Any]          =	"""A black colored car"""
		UpperCamelCase     :Tuple          =	"""A blue colored car"""

		UpperCamelCase     :Tuple          =	torch.manual_seed(0	)
		UpperCamelCase     :List[Any]          =	pipe(
		    prompt=__lowerCamelCase ,  source_prompt=__lowerCamelCase ,  image=__lowerCamelCase ,  num_inference_steps=100 ,  eta=0.1 ,  strength=0.85 ,  guidance_scale=3 ,  source_guidance_scale=1 ,  generator=__lowerCamelCase ,  output_type="""np""" ,  )
		UpperCamelCase     :Optional[int]          =	output.images

		# the values aren't exactly equal, but the images look the same visually
		assert np.abs(image - expected_image	).max() < 5E-1





	def 						_A (		self				:					Optional[Any]	):
		UpperCamelCase     :Dict          =	load_image(
		    """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"""
		    """/cycle-diffusion/black_colored_car.png"""	)
		UpperCamelCase     :Any          =	load_numpy(
		    """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/cycle-diffusion/blue_colored_car.npy"""	)
		UpperCamelCase     :Optional[Any]          =	init_image.resize((512, 512)	)

		UpperCamelCase     :str          =	"""CompVis/stable-diffusion-v1-4"""
		UpperCamelCase     :Union[str, Any]          =	DDIMScheduler.from_pretrained(__lowerCamelCase ,  subfolder="""scheduler"""	)
		UpperCamelCase     :str          =	CycleDiffusionPipeline.from_pretrained(__lowerCamelCase ,  scheduler=__lowerCamelCase ,  safety_checker=__lowerCamelCase	)

		pipe.to(__lowerCamelCase	)
		pipe.set_progress_bar_config(disable=__lowerCamelCase	)
		pipe.enable_attention_slicing()

		UpperCamelCase     :Dict          =	"""A black colored car"""
		UpperCamelCase     :str          =	"""A blue colored car"""

		UpperCamelCase     :List[Any]          =	torch.manual_seed(0	)
		UpperCamelCase     :Optional[int]          =	pipe(
		    prompt=__lowerCamelCase ,  source_prompt=__lowerCamelCase ,  image=__lowerCamelCase ,  num_inference_steps=100 ,  eta=0.1 ,  strength=0.85 ,  guidance_scale=3 ,  source_guidance_scale=1 ,  generator=__lowerCamelCase ,  output_type="""np""" ,  )
		UpperCamelCase     :Optional[int]          =	output.images

		assert np.abs(image - expected_image	).max() < 2E-2


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


__A				=logging.get_logger(__name__)





def 			lowerCamelCase_ (				):
				# Get the sagemaker specific mp parameters from smp_options variable.
				lowerCamelCase_						=			os.getenv("SM_HP_MP_PARAMETERS"	, "{}"    )
				try:
								# Parse it and check the field "partitions" is included, it is required for model parallel.
								lowerCamelCase_						=			json.loads(lowerCamelCase__    )
								if "partitions" not in smp_options:
												return False
				except json.JSONDecodeError:
								return False

				# Get the sagemaker specific framework parameters from mpi_options variable.
				lowerCamelCase_						=			os.getenv("SM_FRAMEWORK_PARAMS"	, "{}"    )
				try:
								# Parse it and check the field "sagemaker_distributed_dataparallel_enabled".
								lowerCamelCase_						=			json.loads(lowerCamelCase__    )
								if not mpi_options.get("sagemaker_mpi_enabled"	, lowerCamelCase__    ):
												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     _SCREAMING_SNAKE_CASE     (      snake_case_						):
	lowerCAmelCase__         = field(
	    default=''			,       metadata={'help': 'Used by the SageMaker launcher to send mp-specific args. Ignored in SageMakerTrainer'}			,       )



	def     SCREAMING_SNAKE_CASE_(    self )				->       Tuple:
					super().__post_init__()
					warnings.warn(
					    "`SageMakerTrainingArguments` is deprecated and will be removed in v5 of Transformers. You can use "
					    "`TrainingArguments` instead."						,  lowercase						,  )



	@cached_property
	def     SCREAMING_SNAKE_CASE_(    self )				->       "torch.device":
					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:
									lowerCamelCase_						=			torch.device("cpu" )
									lowerCamelCase_						=			0
					elif is_sagemaker_model_parallel_available():
									lowerCamelCase_						=			smp.local_rank()
									lowerCamelCase_						=			torch.device("cuda"						,  lowercase )
									lowerCamelCase_						=			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 )
									lowerCamelCase_						=			int(os.getenv("SMDATAPARALLEL_LOCAL_RANK" ) )
									lowerCamelCase_						=			torch.device("cuda"						,  self.local_rank )
									lowerCamelCase_						=			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
									lowerCamelCase_						=			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.
									lowerCamelCase_						=			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 )
									lowerCamelCase_						=			torch.device("cuda"						,  self.local_rank )
									lowerCamelCase_						=			1

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

					return device



	@property
	def     SCREAMING_SNAKE_CASE_(    self )				->       Tuple:
					if is_sagemaker_model_parallel_available():
									return smp.dp_size()

					return super().world_size



	@property
	def     SCREAMING_SNAKE_CASE_(    self )				->       List[str]:
					return not is_sagemaker_model_parallel_available()



	@property
	def     SCREAMING_SNAKE_CASE_(    self )				->       Dict:
					return False


