



"""simple docstring"""


import os
import tempfile
import unittest

import numpy as np

from diffusers.utils import is_flax_available
from diffusers.utils.testing_utils import require_flax, slow


if is_flax_available():
				import jax
				import jax.numpy as jnp
				from flax.jax_utils import replicate
				from flax.training.common_utils import shard

				from diffusers import FlaxDDIMScheduler, FlaxDiffusionPipeline, FlaxStableDiffusionPipeline







@require_flax
class 				A_      (					unittest.TestCase    ):



						"""simple docstring"""






						def      UpperCAmelCase__					( self			:Dict )    ->      str:
												with tempfile.TemporaryDirectory() as tmpdirname:
																		# pipeline has Flax weights
																		UpperCAmelCase					=			FlaxDiffusionPipeline.from_pretrained(
																		    'hf-internal-testing/tiny-stable-diffusion-pipe'       ,  safety_checker=lowercase_       ,  cache_dir=lowercase_ )

																		UpperCAmelCase					=			[t[-1] for t in os.walk(os.path.join(lowercase_       ,  os.listdir(lowercase_ )[0]       ,  'snapshots' ) )]
																		UpperCAmelCase					=			[item for sublist in all_root_files for item in sublist]

																		# None of the downloaded files should be a PyTorch file even if we have some here:
																		# https://huggingface.co/hf-internal-testing/tiny-stable-diffusion-pipe/blob/main/unet/diffusion_pytorch_model.bin
																		assert not any(f.endswith('.bin' ) for f in files )







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



						"""simple docstring"""






						def      UpperCAmelCase__					( self			:int )    ->      List[str]:
												UpperCAmelCase				,   UpperCAmelCase					=			FlaxStableDiffusionPipeline.from_pretrained(
												    'hf-internal-testing/tiny-stable-diffusion-pipe'       ,  safety_checker=lowercase_ )

												UpperCAmelCase					=			(
												    'A cinematic film still of Morgan Freeman starring as Jimi Hendrix, portrait, 40mm lens, shallow depth of'
												    ' field, close up, split lighting, cinematic'
												)

												UpperCAmelCase					=			jax.random.PRNGKey(0 )
												UpperCAmelCase					=			4

												UpperCAmelCase					=			jax.device_count()
												UpperCAmelCase					=			num_samples * [prompt]
												UpperCAmelCase					=			pipeline.prepare_inputs(lowercase_ )

												# shard inputs and rng
												UpperCAmelCase					=			replicate(lowercase_ )
												UpperCAmelCase					=			jax.random.split(lowercase_       ,  lowercase_ )
												UpperCAmelCase					=			shard(lowercase_ )

												UpperCAmelCase					=			pipeline(lowercase_       ,  lowercase_       ,  lowercase_       ,  lowercase_       ,  jit=lowercase_ ).images

												assert images.shape == (num_samples, 1, 64, 64, 3)
												if jax.device_count() == 8:
																		assert np.abs(np.abs(images[0, 0, :2, :2, -2:]       ,  dtype=np.floataa ).sum() - 4.151_4745 ) < 1E-3
																		assert np.abs(np.abs(lowercase_       ,  dtype=np.floataa ).sum() - 4_9947.875 ) < 5E-1

												UpperCAmelCase					=			pipeline.numpy_to_pil(np.asarray(images.reshape((num_samples,) + images.shape[-3:] ) ) )
												assert len(lowercase_ ) == num_samples






						def      UpperCAmelCase__					( self			:List[str] )    ->      Optional[Any]:
												UpperCAmelCase				,   UpperCAmelCase					=			FlaxStableDiffusionPipeline.from_pretrained(
												    'CompVis/stable-diffusion-v1-4'       ,  revision='flax'       ,  safety_checker=lowercase_ )

												UpperCAmelCase					=			(
												    'A cinematic film still of Morgan Freeman starring as Jimi Hendrix, portrait, 40mm lens, shallow depth of'
												    ' field, close up, split lighting, cinematic'
												)

												UpperCAmelCase					=			jax.random.PRNGKey(0 )
												UpperCAmelCase					=			50

												UpperCAmelCase					=			jax.device_count()
												UpperCAmelCase					=			num_samples * [prompt]
												UpperCAmelCase					=			pipeline.prepare_inputs(lowercase_ )

												# shard inputs and rng
												UpperCAmelCase					=			replicate(lowercase_ )
												UpperCAmelCase					=			jax.random.split(lowercase_       ,  lowercase_ )
												UpperCAmelCase					=			shard(lowercase_ )

												UpperCAmelCase					=			pipeline(lowercase_       ,  lowercase_       ,  lowercase_       ,  lowercase_       ,  jit=lowercase_ ).images

												assert images.shape == (num_samples, 1, 5_12, 5_12, 3)
												if jax.device_count() == 8:
																		assert np.abs((np.abs(images[0, 0, :2, :2, -2:]       ,  dtype=np.floataa ).sum() - 0.0565_2401) ) < 1E-3
																		assert np.abs((np.abs(lowercase_       ,  dtype=np.floataa ).sum() - 238_3808.2) ) < 5E-1






						def      UpperCAmelCase__					( self			:str )    ->      Optional[Any]:
												UpperCAmelCase				,   UpperCAmelCase					=			FlaxStableDiffusionPipeline.from_pretrained(
												    'CompVis/stable-diffusion-v1-4'       ,  revision='bf16'       ,  dtype=jnp.bfloataa       ,  safety_checker=lowercase_ )

												UpperCAmelCase					=			(
												    'A cinematic film still of Morgan Freeman starring as Jimi Hendrix, portrait, 40mm lens, shallow depth of'
												    ' field, close up, split lighting, cinematic'
												)

												UpperCAmelCase					=			jax.random.PRNGKey(0 )
												UpperCAmelCase					=			50

												UpperCAmelCase					=			jax.device_count()
												UpperCAmelCase					=			num_samples * [prompt]
												UpperCAmelCase					=			pipeline.prepare_inputs(lowercase_ )

												# shard inputs and rng
												UpperCAmelCase					=			replicate(lowercase_ )
												UpperCAmelCase					=			jax.random.split(lowercase_       ,  lowercase_ )
												UpperCAmelCase					=			shard(lowercase_ )

												UpperCAmelCase					=			pipeline(lowercase_       ,  lowercase_       ,  lowercase_       ,  lowercase_       ,  jit=lowercase_ ).images

												assert images.shape == (num_samples, 1, 5_12, 5_12, 3)
												if jax.device_count() == 8:
																		assert np.abs((np.abs(images[0, 0, :2, :2, -2:]       ,  dtype=np.floataa ).sum() - 0.0400_3906) ) < 1E-3
																		assert np.abs((np.abs(lowercase_       ,  dtype=np.floataa ).sum() - 237_3516.75) ) < 5E-1






						def      UpperCAmelCase__					( self			:Tuple )    ->      List[str]:
												UpperCAmelCase				,   UpperCAmelCase					=			FlaxStableDiffusionPipeline.from_pretrained(
												    'CompVis/stable-diffusion-v1-4'       ,  revision='bf16'       ,  dtype=jnp.bfloataa )

												UpperCAmelCase					=			(
												    'A cinematic film still of Morgan Freeman starring as Jimi Hendrix, portrait, 40mm lens, shallow depth of'
												    ' field, close up, split lighting, cinematic'
												)

												UpperCAmelCase					=			jax.random.PRNGKey(0 )
												UpperCAmelCase					=			50

												UpperCAmelCase					=			jax.device_count()
												UpperCAmelCase					=			num_samples * [prompt]
												UpperCAmelCase					=			pipeline.prepare_inputs(lowercase_ )

												# shard inputs and rng
												UpperCAmelCase					=			replicate(lowercase_ )
												UpperCAmelCase					=			jax.random.split(lowercase_       ,  lowercase_ )
												UpperCAmelCase					=			shard(lowercase_ )

												UpperCAmelCase					=			pipeline(lowercase_       ,  lowercase_       ,  lowercase_       ,  lowercase_       ,  jit=lowercase_ ).images

												assert images.shape == (num_samples, 1, 5_12, 5_12, 3)
												if jax.device_count() == 8:
																		assert np.abs((np.abs(images[0, 0, :2, :2, -2:]       ,  dtype=np.floataa ).sum() - 0.0400_3906) ) < 1E-3
																		assert np.abs((np.abs(lowercase_       ,  dtype=np.floataa ).sum() - 237_3516.75) ) < 5E-1






						def      UpperCAmelCase__					( self			:int )    ->      str:
												UpperCAmelCase					=			FlaxDDIMScheduler(
												    beta_start=0.0_0085       ,  beta_end=0.012       ,  beta_schedule='scaled_linear'       ,  set_alpha_to_one=lowercase_       ,  steps_offset=1       ,  )

												UpperCAmelCase				,   UpperCAmelCase					=			FlaxStableDiffusionPipeline.from_pretrained(
												    'CompVis/stable-diffusion-v1-4'       ,  revision='bf16'       ,  dtype=jnp.bfloataa       ,  scheduler=lowercase_       ,  safety_checker=lowercase_       ,  )
												UpperCAmelCase					=			scheduler.create_state()

												UpperCAmelCase					=			scheduler_state

												UpperCAmelCase					=			(
												    'A cinematic film still of Morgan Freeman starring as Jimi Hendrix, portrait, 40mm lens, shallow depth of'
												    ' field, close up, split lighting, cinematic'
												)

												UpperCAmelCase					=			jax.random.PRNGKey(0 )
												UpperCAmelCase					=			50

												UpperCAmelCase					=			jax.device_count()
												UpperCAmelCase					=			num_samples * [prompt]
												UpperCAmelCase					=			pipeline.prepare_inputs(lowercase_ )

												# shard inputs and rng
												UpperCAmelCase					=			replicate(lowercase_ )
												UpperCAmelCase					=			jax.random.split(lowercase_       ,  lowercase_ )
												UpperCAmelCase					=			shard(lowercase_ )

												UpperCAmelCase					=			pipeline(lowercase_       ,  lowercase_       ,  lowercase_       ,  lowercase_       ,  jit=lowercase_ ).images

												assert images.shape == (num_samples, 1, 5_12, 5_12, 3)
												if jax.device_count() == 8:
																		assert np.abs((np.abs(images[0, 0, :2, :2, -2:]       ,  dtype=np.floataa ).sum() - 0.0_4504_3945) ) < 1E-3
																		assert np.abs((np.abs(lowercase_       ,  dtype=np.floataa ).sum() - 234_7693.5) ) < 5E-1






						def      UpperCAmelCase__					( self			:Optional[Any] )    ->      List[Any]:
												UpperCAmelCase					=			(
												    'A cinematic film still of Morgan Freeman starring as Jimi Hendrix, portrait, 40mm lens, shallow depth of'
												    ' field, close up, split lighting, cinematic'
												)

												UpperCAmelCase					=			jax.device_count()
												UpperCAmelCase					=			num_samples * [prompt]
												UpperCAmelCase					=			jax.random.split(jax.random.PRNGKey(0 )       ,  lowercase_ )

												UpperCAmelCase				,   UpperCAmelCase					=			FlaxStableDiffusionPipeline.from_pretrained(
												    'CompVis/stable-diffusion-v1-4'       ,  revision='bf16'       ,  dtype=jnp.bfloataa       ,  safety_checker=lowercase_       ,  )

												UpperCAmelCase					=			replicate(lowercase_ )
												UpperCAmelCase					=			pipeline.prepare_inputs(lowercase_ )
												UpperCAmelCase					=			shard(lowercase_ )
												UpperCAmelCase					=			pipeline(lowercase_       ,  lowercase_       ,  lowercase_       ,  jit=lowercase_ ).images
												assert images.shape == (num_samples, 1, 5_12, 5_12, 3)
												UpperCAmelCase					=			images[2, 0, 2_56, 10:17, 1]

												# With memory efficient attention
												UpperCAmelCase				,   UpperCAmelCase					=			FlaxStableDiffusionPipeline.from_pretrained(
												    'CompVis/stable-diffusion-v1-4'       ,  revision='bf16'       ,  dtype=jnp.bfloataa       ,  safety_checker=lowercase_       ,  use_memory_efficient_attention=lowercase_       ,  )

												UpperCAmelCase					=			replicate(lowercase_ )
												UpperCAmelCase					=			pipeline.prepare_inputs(lowercase_ )
												UpperCAmelCase					=			shard(lowercase_ )
												UpperCAmelCase					=			pipeline(lowercase_       ,  lowercase_       ,  lowercase_       ,  jit=lowercase_ ).images
												assert images_eff.shape == (num_samples, 1, 5_12, 5_12, 3)
												UpperCAmelCase					=			images[2, 0, 2_56, 10:17, 1]

												# I checked the results visually and they are very similar. However, I saw that the max diff is `1` and the `sum`
												# over the 8 images is exactly `256`, which is very suspicious. Testing a random slice for now.
												assert abs(slice_eff - slice ).max() < 1E-2










"""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				(       ):
						UpperCAmelCase					=			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				(       ):
						UpperCAmelCase					=			parse_args()

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

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





