






'''simple docstring'''

import random
import unittest

import numpy as np

from transformers.testing_utils import require_torch, require_vision
from transformers.utils import is_torch_available, is_vision_available

from ...test_image_processing_common import ImageProcessingSavingTestMixin, prepare_image_inputs


if is_torch_available():
						import torch

if is_vision_available():
						import PIL

						from transformers import FlavaImageProcessor
						from transformers.image_utils import PILImageResampling
						from transformers.models.flava.image_processing_flava import (
						    FLAVA_CODEBOOK_MEAN,
						    FLAVA_CODEBOOK_STD,
						    FLAVA_IMAGE_MEAN,
						    FLAVA_IMAGE_STD,
						)
else:
						A:   List[str]						=					None






class 				__magic_name__  (						unittest.TestCase					):







	"""simple docstring"""



	def __init__(				self	,   _lowercase	,   _lowercase=7	,   _lowercase=3	,   _lowercase=30	,   _lowercase=400	,   _lowercase=True	,   _lowercase=None	,   _lowercase=True	,   _lowercase=None	,   _lowercase=None	,   _lowercase=True	,   _lowercase=1 / 255	,   _lowercase=True	,   _lowercase=FLAVA_IMAGE_MEAN	,   _lowercase=FLAVA_IMAGE_STD	,   _lowercase=14	,   _lowercase=75	,   _lowercase=None	,   _lowercase=16	,   _lowercase=0.3	,   _lowercase=None	,   _lowercase=True	,   _lowercase=None	,   _lowercase=None	,   _lowercase=True	,   _lowercase=None	,   _lowercase=True	,   _lowercase=True	,   _lowercase=FLAVA_CODEBOOK_MEAN	,   _lowercase=FLAVA_CODEBOOK_STD	,   )    ->	int:
							lowercase_   : Optional[int]   	=      size if size is not None else {'height': 224, 'width': 224}
							lowercase_   : Dict   	=      crop_size if crop_size is not None else {'height': 224, 'width': 224}
							lowercase_   : Any   	=      codebook_size if codebook_size is not None else {'height': 112, 'width': 112}
							lowercase_   : Any   	=      codebook_crop_size if codebook_crop_size is not None else {'height': 112, 'width': 112}

							lowercase_   : Optional[Any]   	=      parent
							lowercase_   : List[Any]   	=      batch_size
							lowercase_   : Optional[Any]   	=      num_channels
							lowercase_   : Dict   	=      do_resize
							lowercase_   : int   	=      do_rescale
							lowercase_   : List[Any]   	=      rescale_factor
							lowercase_   : Tuple   	=      min_resolution
							lowercase_   : Optional[Any]   	=      max_resolution
							lowercase_   : Dict   	=      size
							lowercase_   : str   	=      resample if resample is not None else PILImageResampling.BICUBIC
							lowercase_   : Any   	=      do_normalize
							lowercase_   : Optional[int]   	=      image_mean
							lowercase_   : int   	=      image_std
							lowercase_   : Optional[Any]   	=      do_center_crop
							lowercase_   : Tuple   	=      crop_size

							lowercase_   : str   	=      input_size_patches
							lowercase_   : int   	=      total_mask_patches
							lowercase_   : Union[str, Any]   	=      mask_group_max_patches
							lowercase_   : List[Any]   	=      mask_group_min_patches
							lowercase_   : Union[str, Any]   	=      mask_group_min_aspect_ratio
							lowercase_   : Union[str, Any]   	=      mask_group_max_aspect_ratio

							lowercase_   : Optional[int]   	=      codebook_do_resize
							lowercase_   : Optional[int]   	=      codebook_size
							lowercase_   : Optional[Any]   	=      codebook_resample if codebook_resample is not None else PILImageResampling.LANCZOS
							lowercase_   : Union[str, Any]   	=      codebook_do_center_crop
							lowercase_   : Union[str, Any]   	=      codebook_crop_size
							lowercase_   : Union[str, Any]   	=      codebook_do_map_pixels
							lowercase_   : Dict   	=      codebook_do_normalize
							lowercase_   : Optional[Any]   	=      codebook_image_mean
							lowercase_   : Optional[int]   	=      codebook_image_std
	def 							lowerCamelCase__		(				self							)    ->	str:
							return {
							    "image_mean": self.image_mean,
							    "image_std": self.image_std,
							    "do_normalize": self.do_normalize,
							    "do_resize": self.do_resize,
							    "size": self.size,
							    "resample": self.resample,
							    "do_rescale": self.do_rescale,
							    "rescale_factor": self.rescale_factor,
							    "do_center_crop": self.do_center_crop,
							    "crop_size": self.crop_size,
							    "input_size_patches": self.input_size_patches,
							    "total_mask_patches": self.total_mask_patches,
							    "mask_group_max_patches": self.mask_group_max_patches,
							    "mask_group_min_patches": self.mask_group_min_patches,
							    "mask_group_min_aspect_ratio": self.mask_group_min_aspect_ratio,
							    "mask_group_max_aspect_ratio": self.mask_group_min_aspect_ratio,
							    "codebook_do_resize": self.codebook_do_resize,
							    "codebook_size": self.codebook_size,
							    "codebook_resample": self.codebook_resample,
							    "codebook_do_center_crop": self.codebook_do_center_crop,
							    "codebook_crop_size": self.codebook_crop_size,
							    "codebook_do_map_pixels": self.codebook_do_map_pixels,
							    "codebook_do_normalize": self.codebook_do_normalize,
							    "codebook_image_mean": self.codebook_image_mean,
							    "codebook_image_std": self.codebook_image_std,
							}
	def 							lowerCamelCase__		(				self							)    ->	str:
							return (self.size["height"], self.size["width"])
	def 							lowerCamelCase__		(				self							)    ->	List[str]:
							return (
							    (self.input_size_patches, self.input_size_patches)
							    if not isinstance(self.input_size_patches	,   _lowercase							)
							    else self.input_size_patches
							)







	def 							lowerCamelCase__		(				self							)    ->	Tuple:
							return (self.codebook_size["height"], self.codebook_size["width"])






@require_torch
@require_vision
class 				__magic_name__  (						UpperCAmelCase_,    unittest.TestCase					):







	"""simple docstring"""



	SCREAMING_SNAKE_CASE_       :   Tuple   =	FlavaImageProcessor if is_vision_available() else None
	SCREAMING_SNAKE_CASE_       :   Union[str, Any]   =	None
	def 							lowerCamelCase__		(				self							)    ->	str:
							lowercase_   : Tuple   	=      FlavaImageProcessingTester(self							)
	@property
	def 							lowerCamelCase__		(				self							)    ->	List[Any]:
							return self.image_processor_tester.prepare_image_processor_dict()
	def 							lowerCamelCase__		(				self							)    ->	Optional[int]:
							lowercase_   : int   	=      self.image_processing_class(**self.image_processor_dict							)
							self.assertTrue(hasattr(_lowercase	,   'image_mean'							)							)
							self.assertTrue(hasattr(_lowercase	,   'image_std'							)							)
							self.assertTrue(hasattr(_lowercase	,   'do_normalize'							)							)
							self.assertTrue(hasattr(_lowercase	,   'do_resize'							)							)
							self.assertTrue(hasattr(_lowercase	,   'resample'							)							)
							self.assertTrue(hasattr(_lowercase	,   'crop_size'							)							)
							self.assertTrue(hasattr(_lowercase	,   'do_center_crop'							)							)
							self.assertTrue(hasattr(_lowercase	,   'do_rescale'							)							)
							self.assertTrue(hasattr(_lowercase	,   'rescale_factor'							)							)
							self.assertTrue(hasattr(_lowercase	,   'masking_generator'							)							)
							self.assertTrue(hasattr(_lowercase	,   'codebook_do_resize'							)							)
							self.assertTrue(hasattr(_lowercase	,   'codebook_size'							)							)
							self.assertTrue(hasattr(_lowercase	,   'codebook_resample'							)							)
							self.assertTrue(hasattr(_lowercase	,   'codebook_do_center_crop'							)							)
							self.assertTrue(hasattr(_lowercase	,   'codebook_crop_size'							)							)
							self.assertTrue(hasattr(_lowercase	,   'codebook_do_map_pixels'							)							)
							self.assertTrue(hasattr(_lowercase	,   'codebook_do_normalize'							)							)
							self.assertTrue(hasattr(_lowercase	,   'codebook_image_mean'							)							)
							self.assertTrue(hasattr(_lowercase	,   'codebook_image_std'							)							)
	def 							lowerCamelCase__		(				self							)    ->	List[str]:
							lowercase_   : Tuple   	=      self.image_processing_class.from_dict(self.image_processor_dict							)
							self.assertEqual(image_processor.size	,   {'height': 224, 'width': 224}							)
							self.assertEqual(image_processor.crop_size	,   {'height': 224, 'width': 224}							)
							self.assertEqual(image_processor.codebook_size	,   {'height': 112, 'width': 112}							)
							self.assertEqual(image_processor.codebook_crop_size	,   {'height': 112, 'width': 112}							)

							lowercase_   : Any   	=      self.image_processing_class.from_dict(
							    self.image_processor_dict	,   size=42	,   crop_size=84	,   codebook_size=33	,   codebook_crop_size=66							)
							self.assertEqual(image_processor.size	,   {'height': 42, 'width': 42}							)
							self.assertEqual(image_processor.crop_size	,   {'height': 84, 'width': 84}							)
							self.assertEqual(image_processor.codebook_size	,   {'height': 33, 'width': 33}							)
							self.assertEqual(image_processor.codebook_crop_size	,   {'height': 66, 'width': 66}							)
	def 							lowerCamelCase__		(				self							)    ->	Dict:
							pass
	def 							lowerCamelCase__		(				self							)    ->	Optional[Any]:
							# Initialize image_processing
							lowercase_   : List[Any]   	=      self.image_processing_class(**self.image_processor_dict							)
							# create random PIL images
							lowercase_   : str   	=      prepare_image_inputs(self.image_processor_tester	,   equal_resolution=_lowercase							)
							for image in image_inputs:
													self.assertIsInstance(_lowercase	,   PIL.Image.Image							)

							# Test not batched input
							lowercase_   : List[str]   	=      image_processing(image_inputs[0]	,   return_tensors='pt'							)

							# Test no bool masked pos
							self.assertFalse('bool_masked_pos' in encoded_images							)

							lowercase_	,   lowercase_   : Optional[int]   	=      self.image_processor_tester.get_expected_image_size()

							self.assertEqual(
							    encoded_images.pixel_values.shape	,   (1, self.image_processor_tester.num_channels, expected_height, expected_width)	,   )

							# Test batched
							lowercase_   : List[Any]   	=      image_processing(_lowercase	,   return_tensors='pt'							)
							lowercase_	,   lowercase_   : int   	=      self.image_processor_tester.get_expected_image_size()

							# Test no bool masked pos
							self.assertFalse('bool_masked_pos' in encoded_images							)

							self.assertEqual(
							    encoded_images.pixel_values.shape	,   (
							        self.image_processor_tester.batch_size,
							        self.image_processor_tester.num_channels,
							        expected_height,
							        expected_width,
							    )	,   )
	def 							lowerCamelCase__		(				self	,   _lowercase	,   _lowercase							)    ->	List[Any]:
							# Initialize image_processing
							lowercase_   : List[str]   	=      self.image_processing_class(**self.image_processor_dict							)
							# create random tensors
							lowercase_   : Dict   	=      prepare_image_inputs(self.image_processor_tester	,   equal_resolution=_lowercase	,   **_lowercase							)
							for image in image_inputs:
													self.assertIsInstance(_lowercase	,   _lowercase							)

							# Test not batched input
							lowercase_   : Optional[Any]   	=      image_processing(image_inputs[0]	,   return_tensors='pt'							)

							lowercase_	,   lowercase_   : Optional[int]   	=      self.image_processor_tester.get_expected_image_size()
							self.assertEqual(
							    encoded_images.pixel_values.shape	,   (1, self.image_processor_tester.num_channels, expected_height, expected_width)	,   )

							lowercase_   : Union[str, Any]   	=      image_processing(_lowercase	,   return_image_mask=_lowercase	,   return_tensors='pt'							)

							lowercase_	,   lowercase_   : List[str]   	=      self.image_processor_tester.get_expected_image_size()
							self.assertEqual(
							    encoded_images.pixel_values.shape	,   (
							        self.image_processor_tester.batch_size,
							        self.image_processor_tester.num_channels,
							        expected_height,
							        expected_width,
							    )	,   )

							lowercase_	,   lowercase_   : Any   	=      self.image_processor_tester.get_expected_mask_size()
							self.assertEqual(
							    encoded_images.bool_masked_pos.shape	,   (
							        self.image_processor_tester.batch_size,
							        expected_height,
							        expected_width,
							    )	,   )

							# Test batched
							lowercase_   : Optional[int]   	=      image_processing(_lowercase	,   return_tensors='pt'							).pixel_values

							lowercase_	,   lowercase_   : Optional[Any]   	=      self.image_processor_tester.get_expected_image_size()
							self.assertEqual(
							    encoded_images.shape	,   (
							        self.image_processor_tester.batch_size,
							        self.image_processor_tester.num_channels,
							        expected_height,
							        expected_width,
							    )	,   )

							# Test masking
							lowercase_   : Optional[Any]   	=      image_processing(_lowercase	,   return_image_mask=_lowercase	,   return_tensors='pt'							)

							lowercase_	,   lowercase_   : List[str]   	=      self.image_processor_tester.get_expected_image_size()
							self.assertEqual(
							    encoded_images.pixel_values.shape	,   (
							        self.image_processor_tester.batch_size,
							        self.image_processor_tester.num_channels,
							        expected_height,
							        expected_width,
							    )	,   )

							lowercase_	,   lowercase_   : Optional[Any]   	=      self.image_processor_tester.get_expected_mask_size()
							self.assertEqual(
							    encoded_images.bool_masked_pos.shape	,   (
							        self.image_processor_tester.batch_size,
							        expected_height,
							        expected_width,
							    )	,   )
	def 							lowerCamelCase__		(				self							)    ->	Tuple:
							self._test_call_framework(np.ndarray	,   prepare_kwargs={'numpify': True}							)
	def 							lowerCamelCase__		(				self							)    ->	Tuple:
							self._test_call_framework(torch.Tensor	,   prepare_kwargs={'torchify': True}							)
	def 							lowerCamelCase__		(				self							)    ->	List[Any]:
							# Initialize image_processing
							random.seed(1234							)
							lowercase_   : Optional[Any]   	=      self.image_processing_class(**self.image_processor_dict							)
							lowercase_   : Union[str, Any]   	=      prepare_image_inputs(self.image_processor_tester	,   equal_resolution=_lowercase	,   torchify=_lowercase							)

							# Test not batched input
							lowercase_   : Tuple   	=      image_processing(image_inputs[0]	,   return_image_mask=_lowercase	,   return_tensors='pt'							)
							self.assertEqual(encoded_images.bool_masked_pos.sum().item()	,   75							)







	def 							lowerCamelCase__		(				self							)    ->	Any:
							# Initialize image_processing
							lowercase_   : Any   	=      self.image_processing_class(**self.image_processor_dict							)
							# create random PIL images
							lowercase_   : str   	=      prepare_image_inputs(self.image_processor_tester	,   equal_resolution=_lowercase							)
							for image in image_inputs:
													self.assertIsInstance(_lowercase	,   PIL.Image.Image							)

							# Test not batched input
							lowercase_   : List[str]   	=      image_processing(image_inputs[0]	,   return_codebook_pixels=_lowercase	,   return_tensors='pt'							)
							lowercase_	,   lowercase_   : Union[str, Any]   	=      self.image_processor_tester.get_expected_codebook_image_size()
							self.assertEqual(
							    encoded_images.codebook_pixel_values.shape	,   (1, self.image_processor_tester.num_channels, expected_height, expected_width)	,   )

							# Test batched
							lowercase_   : Optional[int]   	=      image_processing(_lowercase	,   return_codebook_pixels=_lowercase	,   return_tensors='pt'							)
							lowercase_	,   lowercase_   : int   	=      self.image_processor_tester.get_expected_codebook_image_size()
							self.assertEqual(
							    encoded_images.codebook_pixel_values.shape	,   (
							        self.image_processor_tester.batch_size,
							        self.image_processor_tester.num_channels,
							        expected_height,
							        expected_width,
							    )	,   )








'''simple docstring'''

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

import torch_xla.distributed.xla_multiprocessing as xmp






def 						_UpperCAmelCase     (	)      ->      Optional[Any]:





						"""simple docstring"""
						lowercase_   : 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=a    ,      default=1    ,      help='Number of TPU cores to use (1 or 8).'				)

						# positional
						parser.add_argument(
						    'training_script'    ,      type=a    ,      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=a				)

						return parser.parse_args()






def 						_UpperCAmelCase     (	)      ->      str:





						"""simple docstring"""
						lowercase_   : Optional[int]   	=      parse_args()

						# Import training_script as a module.
						lowercase_   : str   	=      Path(args.training_script				)
						sys.path.append(str(script_fpath.parent.resolve()				)				)
						lowercase_   : int   	=      script_fpath.stem
						lowercase_   : List[Any]   	=      importlib.import_module(a				)

						# Patch sys.argv
						lowercase_   : List[str]   	=      [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()
