






import random
import unittest

import torch

from diffusers import IFImgaImgSuperResolutionPipeline
from diffusers.utils import floats_tensor
from diffusers.utils.import_utils import is_xformers_available
from diffusers.utils.testing_utils import skip_mps, torch_device

from ..pipeline_params import TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS, TEXT_GUIDED_IMAGE_VARIATION_PARAMS
from ..test_pipelines_common import PipelineTesterMixin
from . import IFPipelineTesterMixin

@skip_mps
class     _SCREAMING_SNAKE_CASE       (			UpperCamelCase_		,   UpperCamelCase_		,   unittest.TestCase      ):
	snake_case__							:     Tuple     	=  IFImgaImgSuperResolutionPipeline
	snake_case__							:     List[Any]     	=  TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {"""width""", """height"""}
	snake_case__							:     Optional[int]     	=  TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS.union({"""original_image"""}      )
	snake_case__							:     Any     	=  PipelineTesterMixin.required_optional_params - {"""latents"""}





	def 						_A (		self				:					str	):
		return self._get_superresolution_dummy_components()





	def 						_A (		self				:					str ,  __lowerCamelCase				:					Any ,  __lowerCamelCase				:					Optional[int]=0	):
		if str(A_	).startswith("""mps"""	):
			UpperCamelCase     :Dict          =	torch.manual_seed(A_	)
		else:
			UpperCamelCase     :Optional[Any]          =	torch.Generator(device=A_	).manual_seed(A_	)

		UpperCamelCase     :Any          =	floats_tensor((1, 3, 32, 32) ,  rng=random.Random(A_	)	).to(A_	)
		UpperCamelCase     :Optional[int]          =	floats_tensor((1, 3, 16, 16) ,  rng=random.Random(A_	)	).to(A_	)

		UpperCamelCase     :Any          =	{
		    '''prompt''': '''A painting of a squirrel eating a burger''',
		    '''image''': image,
		    '''original_image''': original_image,
		    '''generator''': generator,
		    '''num_inference_steps''': 2,
		    '''output_type''': '''numpy''',
		}

		return inputs





	@unittest.skipIf(
	    torch_device != """cuda""" or not is_xformers_available() ,  reason="""XFormers attention is only available with CUDA and `xformers` installed""" ,  )
	def 						_A (		self				:					Tuple	):
		self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=1E-3	)





	def 						_A (		self				:					List[Any]	):
		self._test_save_load_optional_components()





	@unittest.skipIf(torch_device != """cuda""" ,  reason="""float16 requires CUDA"""	)
	def 						_A (		self				:					Dict	):
		# Due to non-determinism in save load of the hf-internal-testing/tiny-random-t5 text encoder
		super().test_save_load_floataa(expected_max_diff=1E-1	)





	def 						_A (		self				:					str	):
		self._test_attention_slicing_forward_pass(expected_max_diff=1E-2	)





	def 						_A (		self				:					List[str]	):
		self._test_save_load_local()





	def 						_A (		self				:					List[str]	):
		self._test_inference_batch_single_identical(
		    expected_max_diff=1E-2 ,  )





import importlib
import sys
from argparse import REMAINDER, ArgumentParser
from pathlib import Path

import torch_xla.distributed.xla_multiprocessing as xmp


def 			UpperCamelCase(  ):
		lowerCAmelCase_			:     List[str]									= ArgumentParser(
		    description=(
		        '''PyTorch TPU distributed training launch '''
		        '''helper utility that will spawn up '''
		        '''multiple distributed processes'''
		    ) )

		# Optional arguments for the launch helper
		parser.add_argument('''--num_cores'''     ,type=__UpperCamelCase     ,default=1     ,help='''Number of TPU cores to use (1 or 8).''' )

		# positional
		parser.add_argument(
		    '''training_script'''     ,type=__UpperCamelCase     ,help=(
		        '''The full path to the single TPU training '''
		        '''program/script to be launched in parallel, '''
		        '''followed by all the arguments for the '''
		        '''training script'''
		    )     ,)

		# rest from the training program
		parser.add_argument('''training_script_args'''     ,nargs=__UpperCamelCase )

		return parser.parse_args()


def 			UpperCamelCase(  ):
		lowerCAmelCase_			:     str									= parse_args()

		# Import training_script as a module.
		lowerCAmelCase_			:     str									= Path(args.training_script )
		sys.path.append(str(script_fpath.parent.resolve() ) )
		lowerCAmelCase_			:     Tuple									= script_fpath.stem
		lowerCAmelCase_			:     Union[str, Any]									= importlib.import_module(__UpperCamelCase )

		# Patch sys.argv
		lowerCAmelCase_			:     Optional[int]									= [args.training_script] + args.training_script_args + ['''--tpu_num_cores''', str(args.num_cores )]
		xmp.spawn(mod._mp_fn     ,args=()     ,nprocs=args.num_cores )


if __name__ == "__main__":
						main()





