






import gc
import random
import unittest

import numpy as np
import torch
from transformers import XLMRobertaTokenizer

from diffusers import (
    AltDiffusionImgaImgPipeline,
    AutoencoderKL,
    PNDMScheduler,
    UNetaDConditionModel,
)
from diffusers.image_processor import VaeImageProcessor
from diffusers.pipelines.alt_diffusion.modeling_roberta_series import (
    RobertaSeriesConfig,
    RobertaSeriesModelWithTransformation,
)
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


enable_full_determinism()

class        lowerCAmelCase_				(      unittest.TestCase	):







				"""simple docstring"""




				def 				__a   ( self     :List[Any]     ):
					# clean up the VRAM after each test
					super().tearDown()
					gc.collect()
					torch.cuda.empty_cache()







				@property
				def 				__a   ( self     :Dict     ):
					UpperCamelCase__ :Optional[int]						=      1
					UpperCamelCase__ :Any						=      3
					UpperCamelCase__ :Optional[Any]						=      (32, 32)

					UpperCamelCase__ :Dict						=      floats_tensor((batch_size, num_channels) + sizes  ,     rng=random.Random(0     )     ).to(lowerCAmelCase__     )
					return image







				@property
				def 				__a   ( self     :List[str]     ):
					torch.manual_seed(0     )
					UpperCamelCase__ :List[Any]						=      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  ,     )
					return model







				@property
				def 				__a   ( self     :str     ):
					torch.manual_seed(0     )
					UpperCamelCase__ :str						=      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  ,     )
					return model







				@property
				def 				__a   ( self     :Optional[Any]     ):
					torch.manual_seed(0     )
					UpperCamelCase__ :List[Any]						=      RobertaSeriesConfig(
					    hidden_size=32  ,     project_dim=32  ,     intermediate_size=37  ,     layer_norm_eps=1e-05  ,     num_attention_heads=4  ,     num_hidden_layers=5  ,     pad_token_id=1  ,     vocab_size=50_06  ,     )
					return RobertaSeriesModelWithTransformation(lowerCAmelCase__     )







				@property
				def 				__a   ( self     :Any     ):
					def extract(*lowerCamelCase__     :str  ,     **lowerCamelCase__     :List[str]     ):

						class        lowerCAmelCase_				:







										"""simple docstring"""




										def __init__( self     :List[str]     ):
											UpperCamelCase__ :Optional[Any]						=      torch.ones([0]     )







										def 				__a   ( self     :Optional[int]  ,     lowerCamelCase__     :Optional[int]     ):
											self.pixel_values.to(lowerCAmelCase__     )
											return self

						return Out()

					return extract







				def 				__a   ( self     :Dict     ):
					UpperCamelCase__ :List[Any]						=      """cpu"""  # ensure determinism for the device-dependent torch.Generator
					UpperCamelCase__ :List[Any]						=      self.dummy_cond_unet
					UpperCamelCase__ :str						=      PNDMScheduler(skip_prk_steps=lowerCAmelCase__     )
					UpperCamelCase__ :Any						=      self.dummy_vae
					UpperCamelCase__ :Dict						=      self.dummy_text_encoder
					UpperCamelCase__ :Dict						=      XLMRobertaTokenizer.from_pretrained("""hf-internal-testing/tiny-xlm-roberta"""     )
					UpperCamelCase__ :List[str]						=      77

					UpperCamelCase__ :str						=      self.dummy_image.to(lowerCAmelCase__     )
					UpperCamelCase__ :Optional[Any]						=      init_image / 2 + 0.5

					# make sure here that pndm scheduler skips prk
					UpperCamelCase__ :Dict						=      AltDiffusionImgaImgPipeline(
					    unet=lowerCAmelCase__  ,     scheduler=lowerCAmelCase__  ,     vae=lowerCAmelCase__  ,     text_encoder=lowerCAmelCase__  ,     tokenizer=lowerCAmelCase__  ,     safety_checker=lowerCAmelCase__  ,     feature_extractor=self.dummy_extractor  ,     )
					UpperCamelCase__ :Dict						=      VaeImageProcessor(vae_scale_factor=alt_pipe.vae_scale_factor  ,     do_normalize=lowerCAmelCase__     )
					UpperCamelCase__ :Union[str, Any]						=      alt_pipe.to(lowerCAmelCase__     )
					alt_pipe.set_progress_bar_config(disable=lowerCAmelCase__     )

					UpperCamelCase__ :Dict						=      """A painting of a squirrel eating a burger"""
					UpperCamelCase__ :Tuple						=      torch.Generator(device=lowerCAmelCase__     ).manual_seed(0     )
					UpperCamelCase__ :Dict						=      alt_pipe(
					    [prompt]  ,     generator=lowerCAmelCase__  ,     guidance_scale=6.0  ,     num_inference_steps=2  ,     output_type="""np"""  ,     image=lowerCAmelCase__  ,     )

					UpperCamelCase__ :Union[str, Any]						=      output.images

					UpperCamelCase__ :Union[str, Any]						=      torch.Generator(device=lowerCAmelCase__     ).manual_seed(0     )
					UpperCamelCase__ :List[str]						=      alt_pipe(
					    [prompt]  ,     generator=lowerCAmelCase__  ,     guidance_scale=6.0  ,     num_inference_steps=2  ,     output_type="""np"""  ,     image=lowerCAmelCase__  ,     return_dict=lowerCAmelCase__  ,     )[0]

					UpperCamelCase__ :int						=      image[0, -3:, -3:, -1]
					UpperCamelCase__ :Optional[int]						=      image_from_tuple[0, -3:, -3:, -1]

					assert image.shape == (1, 32, 32, 3)
					UpperCamelCase__ :Optional[Any]						=      np.array([0.4427, 0.3731, 0.4249, 0.4941, 0.4546, 0.4148, 0.4193, 0.4666, 0.4499]     )

					assert np.abs(image_slice.flatten() - expected_slice     ).max() < 5e-3
					assert np.abs(image_from_tuple_slice.flatten() - expected_slice     ).max() < 5e-3







				@unittest.skipIf(torch_device != """cuda"""  ,     """This test requires a GPU"""     )
				def 				__a   ( self     :str     ):
					UpperCamelCase__ :Optional[Any]						=      self.dummy_cond_unet
					UpperCamelCase__ :Dict						=      PNDMScheduler(skip_prk_steps=lowerCAmelCase__     )
					UpperCamelCase__ :Tuple						=      self.dummy_vae
					UpperCamelCase__ :Dict						=      self.dummy_text_encoder
					UpperCamelCase__ :List[Any]						=      XLMRobertaTokenizer.from_pretrained("""hf-internal-testing/tiny-xlm-roberta"""     )
					UpperCamelCase__ :Optional[int]						=      77

					UpperCamelCase__ :Optional[Any]						=      self.dummy_image.to(lowerCAmelCase__     )

					# put models in fp16
					UpperCamelCase__ :Optional[int]						=      unet.half()
					UpperCamelCase__ :Tuple						=      vae.half()
					UpperCamelCase__ :Optional[int]						=      bert.half()

					# make sure here that pndm scheduler skips prk
					UpperCamelCase__ :Any						=      AltDiffusionImgaImgPipeline(
					    unet=lowerCAmelCase__  ,     scheduler=lowerCAmelCase__  ,     vae=lowerCAmelCase__  ,     text_encoder=lowerCAmelCase__  ,     tokenizer=lowerCAmelCase__  ,     safety_checker=lowerCAmelCase__  ,     feature_extractor=self.dummy_extractor  ,     )
					UpperCamelCase__ :Optional[Any]						=      VaeImageProcessor(vae_scale_factor=alt_pipe.vae_scale_factor  ,     do_normalize=lowerCAmelCase__     )
					UpperCamelCase__ :int						=      alt_pipe.to(lowerCAmelCase__     )
					alt_pipe.set_progress_bar_config(disable=lowerCAmelCase__     )

					UpperCamelCase__ :Any						=      """A painting of a squirrel eating a burger"""
					UpperCamelCase__ :Optional[int]						=      torch.manual_seed(0     )
					UpperCamelCase__ :Tuple						=      alt_pipe(
					    [prompt]  ,     generator=lowerCAmelCase__  ,     num_inference_steps=2  ,     output_type="""np"""  ,     image=lowerCAmelCase__  ,     ).images

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







				@unittest.skipIf(torch_device != """cuda"""  ,     """This test requires a GPU"""     )
				def 				__a   ( self     :Optional[int]     ):
					UpperCamelCase__ :Any						=      load_image(
					    """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"""
					    """/img2img/sketch-mountains-input.jpg"""     )
					# resize to resolution that is divisible by 8 but not 16 or 32
					UpperCamelCase__ :Optional[int]						=      init_image.resize((7_60, 5_04)     )

					UpperCamelCase__ :int						=      """BAAI/AltDiffusion"""
					UpperCamelCase__ :Optional[int]						=      AltDiffusionImgaImgPipeline.from_pretrained(
					    lowerCAmelCase__  ,     safety_checker=lowerCAmelCase__  ,     )
					pipe.to(lowerCAmelCase__     )
					pipe.set_progress_bar_config(disable=lowerCAmelCase__     )
					pipe.enable_attention_slicing()

					UpperCamelCase__ :Tuple						=      """A fantasy landscape, trending on artstation"""

					UpperCamelCase__ :str						=      torch.manual_seed(0     )
					UpperCamelCase__ :Union[str, Any]						=      pipe(
					    prompt=lowerCAmelCase__  ,     image=lowerCAmelCase__  ,     strength=0.75  ,     guidance_scale=7.5  ,     generator=lowerCAmelCase__  ,     output_type="""np"""  ,     )
					UpperCamelCase__ :Any						=      output.images[0]

					UpperCamelCase__ :Union[str, Any]						=      image[2_55:2_58, 3_83:3_86, -1]

					assert image.shape == (5_04, 7_60, 3)
					UpperCamelCase__ :str						=      np.array([0.9358, 0.9397, 0.9599, 0.9901, 1.0000, 1.0000, 0.9882, 1.0000, 1.0000]     )

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

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







				"""simple docstring"""




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







				def 				__a   ( self     :Optional[Any]     ):
					UpperCamelCase__ :int						=      load_image(
					    """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"""
					    """/img2img/sketch-mountains-input.jpg"""     )
					UpperCamelCase__ :Union[str, Any]						=      init_image.resize((7_68, 5_12)     )
					UpperCamelCase__ :int						=      load_numpy(
					    """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/img2img/fantasy_landscape_alt.npy"""     )

					UpperCamelCase__ :Tuple						=      """BAAI/AltDiffusion"""
					UpperCamelCase__ :Any						=      AltDiffusionImgaImgPipeline.from_pretrained(
					    lowerCAmelCase__  ,     safety_checker=lowerCAmelCase__  ,     )
					pipe.to(lowerCAmelCase__     )
					pipe.set_progress_bar_config(disable=lowerCAmelCase__     )
					pipe.enable_attention_slicing()

					UpperCamelCase__ :Any						=      """A fantasy landscape, trending on artstation"""

					UpperCamelCase__ :List[Any]						=      torch.manual_seed(0     )
					UpperCamelCase__ :Optional[int]						=      pipe(
					    prompt=lowerCAmelCase__  ,     image=lowerCAmelCase__  ,     strength=0.75  ,     guidance_scale=7.5  ,     generator=lowerCAmelCase__  ,     output_type="""np"""  ,     )
					UpperCamelCase__ :Union[str, Any]						=      output.images[0]

					assert image.shape == (5_12, 7_68, 3)
					# img2img is flaky across GPUs even in fp32, so using MAE here
					assert np.abs(expected_image - image     ).max() < 1e-2

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

import torch_xla.distributed.xla_multiprocessing as xmp

def  __snake_case							(	):







				'''simple docstring'''


				lowercase         =      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=__magic_name__		,    default=1		,    help="Number of TPU cores to use (1 or 8)."						)

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

				return parser.parse_args()

def  __snake_case							(	):







				'''simple docstring'''


				lowercase         =      parse_args()

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

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