





'''simple docstring'''







from typing import Any, Callable, Dict, List, Optional, Union

import PIL.Image
import torch
from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    DiffusionPipeline,
    LMSDiscreteScheduler,
    PNDMScheduler,
    StableDiffusionImgaImgPipeline,
    StableDiffusionInpaintPipelineLegacy,
    StableDiffusionPipeline,
    UNetaDConditionModel,
)
from diffusers.configuration_utils import FrozenDict
from diffusers.pipelines.stable_diffusion.safety_checker import StableDiffusionSafetyChecker
from diffusers.utils import deprecate, logging


_a					: Optional[int]      =    logging.get_logger(__name__)  # pylint: disable=invalid-name





class     _UpperCAmelCase							(						lowerCAmelCase_   ):
		a   :   List[Any]     =["""safety_checker""", """feature_extractor"""]
		def __init__(  self,__SCREAMING_SNAKE_CASE,__SCREAMING_SNAKE_CASE,__SCREAMING_SNAKE_CASE,__SCREAMING_SNAKE_CASE,__SCREAMING_SNAKE_CASE,__SCREAMING_SNAKE_CASE,__SCREAMING_SNAKE_CASE,__SCREAMING_SNAKE_CASE = True,):



					'''simple docstring'''

					super().__init__()
					if hasattr(scheduler.config,"""steps_offset"""			) and scheduler.config.steps_offset != 1:
								__lowerCAmelCase											=				(
								    f'The configuration file of this scheduler: {scheduler} is outdated. `steps_offset`'
								    f' should be set to 1 instead of {scheduler.config.steps_offset}. Please make sure '
								    """to update the config accordingly as leaving `steps_offset` might led to incorrect results"""
								    """ in future versions. If you have downloaded this checkpoint from the Hugging Face Hub,"""
								    """ it would be very nice if you could open a Pull request for the `scheduler/scheduler_config.json`"""
								    """ file"""
								)
								deprecate("""steps_offset!=1""","""1.0.0""",__SCREAMING_SNAKE_CASE,standard_warn=__SCREAMING_SNAKE_CASE			)
								__lowerCAmelCase											=				dict(scheduler.config			)
								__lowerCAmelCase											=				1
								__lowerCAmelCase											=				FrozenDict(__SCREAMING_SNAKE_CASE			)

					self.register_modules(
					    vae=__SCREAMING_SNAKE_CASE,text_encoder=__SCREAMING_SNAKE_CASE,tokenizer=__SCREAMING_SNAKE_CASE,unet=__SCREAMING_SNAKE_CASE,scheduler=__SCREAMING_SNAKE_CASE,safety_checker=__SCREAMING_SNAKE_CASE,feature_extractor=__SCREAMING_SNAKE_CASE,)
					self.register_to_config(requires_safety_checker=__SCREAMING_SNAKE_CASE			)
		@property
		def   lowerCamelCase__			(  self			):



					'''simple docstring'''

					return {k: getattr(self,__SCREAMING_SNAKE_CASE			) for k in self.config.keys() if not k.startswith("""_"""			)}
		def   lowerCamelCase__			(  self,__SCREAMING_SNAKE_CASE = "auto"			):



					'''simple docstring'''

					if slice_size == "auto":
								# half the attention head size is usually a good trade-off between
								# speed and memory
								__lowerCAmelCase											=				self.unet.config.attention_head_dim // 2
					self.unet.set_attention_slice(__SCREAMING_SNAKE_CASE			)
		def   lowerCamelCase__			(  self			):



					'''simple docstring'''

					self.enable_attention_slicing(__SCREAMING_SNAKE_CASE			)
		@torch.no_grad()
		def   lowerCamelCase__			(  self,__SCREAMING_SNAKE_CASE,__SCREAMING_SNAKE_CASE,__SCREAMING_SNAKE_CASE,__SCREAMING_SNAKE_CASE = 0.8,__SCREAMING_SNAKE_CASE = 50,__SCREAMING_SNAKE_CASE = 7.5,__SCREAMING_SNAKE_CASE = None,__SCREAMING_SNAKE_CASE = 1,__SCREAMING_SNAKE_CASE = 0.0,__SCREAMING_SNAKE_CASE = None,__SCREAMING_SNAKE_CASE = "pil",__SCREAMING_SNAKE_CASE = True,__SCREAMING_SNAKE_CASE = None,__SCREAMING_SNAKE_CASE = 1,):



					'''simple docstring'''

					return StableDiffusionInpaintPipelineLegacy(**self.components			)(
					    prompt=__SCREAMING_SNAKE_CASE,image=__SCREAMING_SNAKE_CASE,mask_image=__SCREAMING_SNAKE_CASE,strength=__SCREAMING_SNAKE_CASE,num_inference_steps=__SCREAMING_SNAKE_CASE,guidance_scale=__SCREAMING_SNAKE_CASE,negative_prompt=__SCREAMING_SNAKE_CASE,num_images_per_prompt=__SCREAMING_SNAKE_CASE,eta=__SCREAMING_SNAKE_CASE,generator=__SCREAMING_SNAKE_CASE,output_type=__SCREAMING_SNAKE_CASE,return_dict=__SCREAMING_SNAKE_CASE,callback=__SCREAMING_SNAKE_CASE,)
		@torch.no_grad()
		def   lowerCamelCase__			(  self,__SCREAMING_SNAKE_CASE,__SCREAMING_SNAKE_CASE,__SCREAMING_SNAKE_CASE = 0.8,__SCREAMING_SNAKE_CASE = 50,__SCREAMING_SNAKE_CASE = 7.5,__SCREAMING_SNAKE_CASE = None,__SCREAMING_SNAKE_CASE = 1,__SCREAMING_SNAKE_CASE = 0.0,__SCREAMING_SNAKE_CASE = None,__SCREAMING_SNAKE_CASE = "pil",__SCREAMING_SNAKE_CASE = True,__SCREAMING_SNAKE_CASE = None,__SCREAMING_SNAKE_CASE = 1,**__SCREAMING_SNAKE_CASE,):



					'''simple docstring'''

					return StableDiffusionImgaImgPipeline(**self.components			)(
					    prompt=__SCREAMING_SNAKE_CASE,image=__SCREAMING_SNAKE_CASE,strength=__SCREAMING_SNAKE_CASE,num_inference_steps=__SCREAMING_SNAKE_CASE,guidance_scale=__SCREAMING_SNAKE_CASE,negative_prompt=__SCREAMING_SNAKE_CASE,num_images_per_prompt=__SCREAMING_SNAKE_CASE,eta=__SCREAMING_SNAKE_CASE,generator=__SCREAMING_SNAKE_CASE,output_type=__SCREAMING_SNAKE_CASE,return_dict=__SCREAMING_SNAKE_CASE,callback=__SCREAMING_SNAKE_CASE,callback_steps=__SCREAMING_SNAKE_CASE,)







		@torch.no_grad()
		def   lowerCamelCase__			(  self,__SCREAMING_SNAKE_CASE,__SCREAMING_SNAKE_CASE = 5_12,__SCREAMING_SNAKE_CASE = 5_12,__SCREAMING_SNAKE_CASE = 50,__SCREAMING_SNAKE_CASE = 7.5,__SCREAMING_SNAKE_CASE = None,__SCREAMING_SNAKE_CASE = 1,__SCREAMING_SNAKE_CASE = 0.0,__SCREAMING_SNAKE_CASE = None,__SCREAMING_SNAKE_CASE = None,__SCREAMING_SNAKE_CASE = "pil",__SCREAMING_SNAKE_CASE = True,__SCREAMING_SNAKE_CASE = None,__SCREAMING_SNAKE_CASE = 1,):



					'''simple docstring'''

					return StableDiffusionPipeline(**self.components			)(
					    prompt=__SCREAMING_SNAKE_CASE,height=__SCREAMING_SNAKE_CASE,width=__SCREAMING_SNAKE_CASE,num_inference_steps=__SCREAMING_SNAKE_CASE,guidance_scale=__SCREAMING_SNAKE_CASE,negative_prompt=__SCREAMING_SNAKE_CASE,num_images_per_prompt=__SCREAMING_SNAKE_CASE,eta=__SCREAMING_SNAKE_CASE,generator=__SCREAMING_SNAKE_CASE,latents=__SCREAMING_SNAKE_CASE,output_type=__SCREAMING_SNAKE_CASE,return_dict=__SCREAMING_SNAKE_CASE,callback=__SCREAMING_SNAKE_CASE,callback_steps=__SCREAMING_SNAKE_CASE,)








'''simple docstring'''







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

import torch_xla.distributed.xla_multiprocessing as xmp





def    _lowerCAmelCase						(						)					->					List[str]:
			__lowerCAmelCase											=				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=lowercase					,	default=1					,	help="""Number of TPU cores to use (1 or 8)."""	)

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

			return parser.parse_args()





def    _lowerCAmelCase						(						)					->					Any:
			__lowerCAmelCase											=				parse_args()

			# Import training_script as a module.
			__lowerCAmelCase											=				Path(args.training_script	)
			sys.path.append(str(script_fpath.parent.resolve()	)	)
			__lowerCAmelCase											=				script_fpath.stem
			__lowerCAmelCase											=				importlib.import_module(lowercase	)

			# Patch sys.argv
			__lowerCAmelCase											=				[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()

