
import gc
import unittest

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

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    EulerAncestralDiscreteScheduler,
    PNDMScheduler,
    StableDiffusionModelEditingPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu, skip_mps

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineKarrasSchedulerTesterMixin, PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()







@skip_mps
class 	A_       (      __lowerCamelCase				,     __lowerCamelCase				,     __lowerCamelCase				,     unittest.TestCase       ):




	'''simple docstring'''


	_UpperCamelCase					:					Dict										=	StableDiffusionModelEditingPipeline
	_UpperCamelCase					:					Tuple										=	TEXT_TO_IMAGE_PARAMS
	_UpperCamelCase					:					List[Any]										=	TEXT_TO_IMAGE_BATCH_PARAMS
	_UpperCamelCase					:					Optional[Any]										=	TEXT_TO_IMAGE_IMAGE_PARAMS
	_UpperCamelCase					:					Tuple										=	TEXT_TO_IMAGE_IMAGE_PARAMS





	def      SCREAMING_SNAKE_CASE__						(	self	):
		torch.manual_seed(0	)
		lowercase					  =   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		,   )
		lowercase					  =   DDIMScheduler()
		torch.manual_seed(0	)
		lowercase					  =   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	)
		lowercase					  =   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=1000		,   )
		lowercase					  =   CLIPTextModel(snake_case	)
		lowercase					  =   CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'	)

		lowercase					  =   {
		    'unet': unet,
		    'scheduler': scheduler,
		    'vae': vae,
		    'text_encoder': text_encoder,
		    'tokenizer': tokenizer,
		    'safety_checker': None,
		    'feature_extractor': None,
		}
		return components





	def      SCREAMING_SNAKE_CASE__						(	self		,   snake_case		,   snake_case=0	):
		lowercase					  =   torch.manual_seed(snake_case	)
		lowercase					  =   {
		    'prompt': 'A field of roses',
		    'generator': generator,
		    # Setting height and width to None to prevent OOMs on CPU.
		    'height': None,
		    'width': None,
		    'num_inference_steps': 2,
		    'guidance_scale': 6.0,
		    'output_type': 'numpy',
		}
		return inputs





	def      SCREAMING_SNAKE_CASE__						(	self	):
		lowercase					  =   'cpu'  # ensure determinism for the device-dependent torch.Generator
		lowercase					  =   self.get_dummy_components()
		lowercase					  =   StableDiffusionModelEditingPipeline(**snake_case	)
		lowercase					  =   sd_pipe.to(snake_case	)
		sd_pipe.set_progress_bar_config(disable=snake_case	)

		lowercase					  =   self.get_dummy_inputs(snake_case	)
		lowercase					  =   sd_pipe(**snake_case	).images
		lowercase					  =   image[0, -3:, -3:, -1]
		assert image.shape == (1, 64, 64, 3)

		lowercase					  =   np.array([0.4_755, 0.5_132, 0.4_976, 0.3_904, 0.3_554, 0.4_765, 0.5_139, 0.5_158, 0.4_889]	)

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





	def      SCREAMING_SNAKE_CASE__						(	self	):
		lowercase					  =   'cpu'  # ensure determinism for the device-dependent torch.Generator
		lowercase					  =   self.get_dummy_components()
		lowercase					  =   StableDiffusionModelEditingPipeline(**snake_case	)
		lowercase					  =   sd_pipe.to(snake_case	)
		sd_pipe.set_progress_bar_config(disable=snake_case	)

		lowercase					  =   self.get_dummy_inputs(snake_case	)
		lowercase					  =   'french fries'
		lowercase					  =   sd_pipe(**snake_case		,   negative_prompt=snake_case	)
		lowercase					  =   output.images
		lowercase					  =   image[0, -3:, -3:, -1]

		assert image.shape == (1, 64, 64, 3)

		lowercase					  =   np.array([0.4_992, 0.5_101, 0.5_004, 0.3_949, 0.3_604, 0.4_735, 0.5_216, 0.5_204, 0.4_913]	)

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





	def      SCREAMING_SNAKE_CASE__						(	self	):
		lowercase					  =   'cpu'  # ensure determinism for the device-dependent torch.Generator
		lowercase					  =   self.get_dummy_components()
		lowercase					  =   EulerAncestralDiscreteScheduler(
		    beta_start=0.00_085		,   beta_end=0.012		,   beta_schedule='scaled_linear'	)
		lowercase					  =   StableDiffusionModelEditingPipeline(**snake_case	)
		lowercase					  =   sd_pipe.to(snake_case	)
		sd_pipe.set_progress_bar_config(disable=snake_case	)

		lowercase					  =   self.get_dummy_inputs(snake_case	)
		lowercase					  =   sd_pipe(**snake_case	).images
		lowercase					  =   image[0, -3:, -3:, -1]

		assert image.shape == (1, 64, 64, 3)

		lowercase					  =   np.array([0.4_747, 0.5_372, 0.4_779, 0.4_982, 0.5_543, 0.4_816, 0.5_238, 0.4_904, 0.5_027]	)

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





	def      SCREAMING_SNAKE_CASE__						(	self	):
		lowercase					  =   'cpu'  # ensure determinism for the device-dependent torch.Generator
		lowercase					  =   self.get_dummy_components()
		lowercase					  =   PNDMScheduler()
		lowercase					  =   StableDiffusionModelEditingPipeline(**snake_case	)
		lowercase					  =   sd_pipe.to(snake_case	)
		sd_pipe.set_progress_bar_config(disable=snake_case	)

		lowercase					  =   self.get_dummy_inputs(snake_case	)
		# the pipeline does not expect pndm so test if it raises error.
		with self.assertRaises(snake_case	):
			lowercase					  =   sd_pipe(**snake_case	).images





	def      SCREAMING_SNAKE_CASE__						(	self	):
		super().test_inference_batch_single_identical(expected_max_diff=5E-3	)





	def      SCREAMING_SNAKE_CASE__						(	self	):
		super().test_attention_slicing_forward_pass(expected_max_diff=5E-3	)







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




	'''simple docstring'''


	def      SCREAMING_SNAKE_CASE__						(	self	):
		super().tearDown()
		gc.collect()
		torch.cuda.empty_cache()





	def      SCREAMING_SNAKE_CASE__						(	self		,   snake_case=0	):
		lowercase					  =   torch.manual_seed(snake_case	)
		lowercase					  =   {
		    'prompt': 'A field of roses',
		    'generator': generator,
		    'num_inference_steps': 3,
		    'guidance_scale': 7.5,
		    'output_type': 'numpy',
		}
		return inputs





	def      SCREAMING_SNAKE_CASE__						(	self	):
		lowercase					  =   'CompVis/stable-diffusion-v1-4'
		lowercase					  =   StableDiffusionModelEditingPipeline.from_pretrained(snake_case		,   safety_checker=snake_case	)
		pipe.to(snake_case	)
		pipe.set_progress_bar_config(disable=snake_case	)
		pipe.enable_attention_slicing()

		lowercase					  =   self.get_inputs()
		lowercase					  =   pipe(**snake_case	).images
		lowercase					  =   image[0, -3:, -3:, -1].flatten()

		assert image.shape == (1, 512, 512, 3)

		lowercase					  =   np.array(
		    [0.6_749_496, 0.6_386_453, 0.51_443_267, 0.66_094_905, 0.61_921_215, 0.5_491_332, 0.5_744_417, 0.58_075_106, 0.5_174_658]	)

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

		# make sure image changes after editing
		pipe.edit_model('A pack of roses'		,   'A pack of blue roses'	)

		lowercase					  =   pipe(**snake_case	).images
		lowercase					  =   image[0, -3:, -3:, -1].flatten()

		assert image.shape == (1, 512, 512, 3)

		assert np.abs(expected_slice - image_slice	).max() > 1E-1





	def      SCREAMING_SNAKE_CASE__						(	self	):
		torch.cuda.empty_cache()
		torch.cuda.reset_max_memory_allocated()
		torch.cuda.reset_peak_memory_stats()

		lowercase					  =   'CompVis/stable-diffusion-v1-4'
		lowercase					  =   DDIMScheduler.from_pretrained(snake_case		,   subfolder='scheduler'	)
		lowercase					  =   StableDiffusionModelEditingPipeline.from_pretrained(
		    snake_case		,   scheduler=snake_case		,   safety_checker=snake_case	)
		lowercase					  =   pipe.to(snake_case	)
		pipe.set_progress_bar_config(disable=snake_case	)
		pipe.enable_attention_slicing(1	)
		pipe.enable_sequential_cpu_offload()

		lowercase					  =   self.get_inputs()
		lowercase					  =   pipe(**snake_case	)

		lowercase					  =   torch.cuda.max_memory_allocated()
		# make sure that less than 4.4 GB is allocated
		assert mem_bytes < 4.4 * 10**9



# This test is meant to be run in on an instance with TPUs like this:
#
#   python examples/pytorch/xla_spawn.py --num_cores=8 tests/test_trainer_tpu.py
#
# Replace 8 with the number of TPU cores you have.
#

import sys
from typing import Dict

from transformers import EvalPrediction, HfArgumentParser, TrainingArguments, is_torch_available
from transformers.utils import logging


UpperCAmelCase     				=			logging.get_logger(__name__)


if is_torch_available():
	import torch
	from torch import nn
	from torch.utils.data import Dataset

	from transformers import Trainer







	class 	A_       (      __lowerCamelCase       ):




		'''simple docstring'''


		def __init__(	self		,   snake_case = 101	):
			lowercase					  =   length





		def __len__(	self	):
			return self.length





		def __getitem__(	self		,   snake_case	):
			return i







	class 	A_       :




		'''simple docstring'''


		def __call__(	self		,   snake_case	):
			return {"input_ids": torch.tensor(snake_case	), "labels": torch.tensor(snake_case	)}







	class 	A_       (      nn.Module       ):




		'''simple docstring'''


		def __init__(	self	):
			super().__init__()
			# Add some (unused) params otherwise DDP will complain.
			lowercase					  =   nn.Linear(120		,   80	)





		def      SCREAMING_SNAKE_CASE__						(	self		,   snake_case		,   snake_case=None	):
			if labels is not None:
				return torch.tensor(0.0		,   device=input_ids.device	), input_ids
			else:
				return input_ids






def 							UpperCAmelCase_		(							):
	lowercase					  =   HfArgumentParser((TrainingArguments,)							)
	sys.argv += ["--output_dir", "./examples"]
	lowercase					  =   parser.parse_args_into_dataclasses()[0]

	logger.warning(
	    F'''Process rank: {training_args.local_rank}, device: {training_args.device}, '''
	    F'''tpu_num_cores: {training_args.tpu_num_cores}'''			,							)

	# Essentially, what we want to verify in the distributed case is
	# that we get all samples back, in the right order.
	# (this is crucial for prediction for instance)
	for dataset_length in [1001, 256, 15]:
		lowercase					  =   DummyDataset(__SCREAMING_SNAKE_CASE							)

		def compute_metrics(__SCREAMING_SNAKE_CASE							) -> Dict:
			lowercase					  =   list(range(len(__SCREAMING_SNAKE_CASE							)							)							)
			lowercase					  =   p.predictions.tolist() == sequential and p.label_ids.tolist() == sequential
			return {"success": success}

		lowercase					  =   Trainer(
		    model=DummyModel()			,							args=__SCREAMING_SNAKE_CASE			,							data_collator=DummyDataCollator()			,							eval_dataset=__SCREAMING_SNAKE_CASE			,							compute_metrics=__SCREAMING_SNAKE_CASE			,							)
		lowercase					  =   trainer.evaluate()
		logger.info(__SCREAMING_SNAKE_CASE							)
		if metrics["eval_success"] is not True:
			logger.error(__SCREAMING_SNAKE_CASE							)
			exit(1							)

		lowercase					  =   trainer.predict(__SCREAMING_SNAKE_CASE							)
		logger.info(p.metrics							)
		if p.metrics["test_success"] is not True:
			logger.error(p.metrics							)
			exit(1							)

		lowercase					  =   2

		lowercase					  =   trainer.evaluate()
		logger.info(__SCREAMING_SNAKE_CASE							)
		if metrics["eval_success"] is not True:
			logger.error(__SCREAMING_SNAKE_CASE							)
			exit(1							)

		lowercase					  =   trainer.predict(__SCREAMING_SNAKE_CASE							)
		logger.info(p.metrics							)
		if p.metrics["test_success"] is not True:
			logger.error(p.metrics							)
			exit(1							)

		lowercase					  =   None

	logger.info('🔥 All distributed tests successful'							)






def 							UpperCAmelCase_		(							__SCREAMING_SNAKE_CASE							):
	# For xla_spawn (TPUs)
	main()


if __name__ == "__main__":
	main()

