


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      ):

							def 				lowerCamelCase_			(							self:     Optional[Any]							)      ->       Optional[int]:





									"""simple docstring"""







									with tempfile.TemporaryDirectory() as tmpdirname:
											# pipeline has Flax weights
											lowercase__						=				FlaxDiffusionPipeline.from_pretrained(
											    '''hf-internal-testing/tiny-stable-diffusion-pipe'''   ,						safety_checker=UpperCamelCase_   ,						cache_dir=UpperCamelCase_							)

											lowercase__						=				[t[-1] for t in os.walk(os.path.join(UpperCamelCase_   ,						os.listdir(UpperCamelCase_							)[0]   ,						'''snapshots'''							)							)]
											lowercase__						=				[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      ):

							def 				lowerCamelCase_			(							self:     List[str]							)      ->       int:





									"""simple docstring"""







									lowercase__					,		lowercase__						=				FlaxStableDiffusionPipeline.from_pretrained(
									    '''hf-internal-testing/tiny-stable-diffusion-pipe'''   ,						safety_checker=UpperCamelCase_							)

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

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

									lowercase__						=				jax.device_count()
									lowercase__						=				num_samples * [prompt]
									lowercase__						=				pipeline.prepare_inputs(UpperCamelCase_							)

									# shard inputs and rng
									lowercase__						=				replicate(UpperCamelCase_							)
									lowercase__						=				jax.random.split(UpperCamelCase_   ,						UpperCamelCase_							)
									lowercase__						=				shard(UpperCamelCase_							)

									lowercase__						=				pipeline(UpperCamelCase_   ,						UpperCamelCase_   ,						UpperCamelCase_   ,						UpperCamelCase_   ,						jit=UpperCamelCase_							).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.1514745							) < 1E-3
											assert np.abs(np.abs(UpperCamelCase_   ,						dtype=np.floataa							).sum() - 49947.875							) < 5E-1

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

							def 				lowerCamelCase_			(							self:     Optional[int]							)      ->       Optional[int]:





									"""simple docstring"""







									lowercase__					,		lowercase__						=				FlaxStableDiffusionPipeline.from_pretrained(
									    '''CompVis/stable-diffusion-v1-4'''   ,						revision='''flax'''   ,						safety_checker=UpperCamelCase_							)

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

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

									lowercase__						=				jax.device_count()
									lowercase__						=				num_samples * [prompt]
									lowercase__						=				pipeline.prepare_inputs(UpperCamelCase_							)

									# shard inputs and rng
									lowercase__						=				replicate(UpperCamelCase_							)
									lowercase__						=				jax.random.split(UpperCamelCase_   ,						UpperCamelCase_							)
									lowercase__						=				shard(UpperCamelCase_							)

									lowercase__						=				pipeline(UpperCamelCase_   ,						UpperCamelCase_   ,						UpperCamelCase_   ,						UpperCamelCase_   ,						jit=UpperCamelCase_							).images

									assert images.shape == (num_samples, 1, 512, 512, 3)
									if jax.device_count() == 8:
											assert np.abs((np.abs(images[0, 0, :2, :2, -2:]   ,						dtype=np.floataa							).sum() - 0.05652401)							) < 1E-3
											assert np.abs((np.abs(UpperCamelCase_   ,						dtype=np.floataa							).sum() - 2383808.2)							) < 5E-1

							def 				lowerCamelCase_			(							self:     Tuple							)      ->       Dict:





									"""simple docstring"""







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

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

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

									lowercase__						=				jax.device_count()
									lowercase__						=				num_samples * [prompt]
									lowercase__						=				pipeline.prepare_inputs(UpperCamelCase_							)

									# shard inputs and rng
									lowercase__						=				replicate(UpperCamelCase_							)
									lowercase__						=				jax.random.split(UpperCamelCase_   ,						UpperCamelCase_							)
									lowercase__						=				shard(UpperCamelCase_							)

									lowercase__						=				pipeline(UpperCamelCase_   ,						UpperCamelCase_   ,						UpperCamelCase_   ,						UpperCamelCase_   ,						jit=UpperCamelCase_							).images

									assert images.shape == (num_samples, 1, 512, 512, 3)
									if jax.device_count() == 8:
											assert np.abs((np.abs(images[0, 0, :2, :2, -2:]   ,						dtype=np.floataa							).sum() - 0.04003906)							) < 1E-3
											assert np.abs((np.abs(UpperCamelCase_   ,						dtype=np.floataa							).sum() - 2373516.75)							) < 5E-1

							def 				lowerCamelCase_			(							self:     Union[str, Any]							)      ->       Any:





									"""simple docstring"""







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

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

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

									lowercase__						=				jax.device_count()
									lowercase__						=				num_samples * [prompt]
									lowercase__						=				pipeline.prepare_inputs(UpperCamelCase_							)

									# shard inputs and rng
									lowercase__						=				replicate(UpperCamelCase_							)
									lowercase__						=				jax.random.split(UpperCamelCase_   ,						UpperCamelCase_							)
									lowercase__						=				shard(UpperCamelCase_							)

									lowercase__						=				pipeline(UpperCamelCase_   ,						UpperCamelCase_   ,						UpperCamelCase_   ,						UpperCamelCase_   ,						jit=UpperCamelCase_							).images

									assert images.shape == (num_samples, 1, 512, 512, 3)
									if jax.device_count() == 8:
											assert np.abs((np.abs(images[0, 0, :2, :2, -2:]   ,						dtype=np.floataa							).sum() - 0.04003906)							) < 1E-3
											assert np.abs((np.abs(UpperCamelCase_   ,						dtype=np.floataa							).sum() - 2373516.75)							) < 5E-1

							def 				lowerCamelCase_			(							self:     Tuple							)      ->       Any:





									"""simple docstring"""







									lowercase__						=				FlaxDDIMScheduler(
									    beta_start=0.00085   ,						beta_end=0.012   ,						beta_schedule='''scaled_linear'''   ,						set_alpha_to_one=UpperCamelCase_   ,						steps_offset=1   ,						)

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

									lowercase__						=				scheduler_state

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

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

									lowercase__						=				jax.device_count()
									lowercase__						=				num_samples * [prompt]
									lowercase__						=				pipeline.prepare_inputs(UpperCamelCase_							)

									# shard inputs and rng
									lowercase__						=				replicate(UpperCamelCase_							)
									lowercase__						=				jax.random.split(UpperCamelCase_   ,						UpperCamelCase_							)
									lowercase__						=				shard(UpperCamelCase_							)

									lowercase__						=				pipeline(UpperCamelCase_   ,						UpperCamelCase_   ,						UpperCamelCase_   ,						UpperCamelCase_   ,						jit=UpperCamelCase_							).images

									assert images.shape == (num_samples, 1, 512, 512, 3)
									if jax.device_count() == 8:
											assert np.abs((np.abs(images[0, 0, :2, :2, -2:]   ,						dtype=np.floataa							).sum() - 0.045043945)							) < 1E-3
											assert np.abs((np.abs(UpperCamelCase_   ,						dtype=np.floataa							).sum() - 2347693.5)							) < 5E-1


							def 				lowerCamelCase_			(							self:     Optional[int]							)      ->       Optional[Any]:





									"""simple docstring"""







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

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

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

									lowercase__						=				replicate(UpperCamelCase_							)
									lowercase__						=				pipeline.prepare_inputs(UpperCamelCase_							)
									lowercase__						=				shard(UpperCamelCase_							)
									lowercase__						=				pipeline(UpperCamelCase_   ,						UpperCamelCase_   ,						UpperCamelCase_   ,						jit=UpperCamelCase_							).images
									assert images.shape == (num_samples, 1, 512, 512, 3)
									lowercase__						=				images[2, 0, 256, 10:17, 1]

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

									lowercase__						=				replicate(UpperCamelCase_							)
									lowercase__						=				pipeline.prepare_inputs(UpperCamelCase_							)
									lowercase__						=				shard(UpperCamelCase_							)
									lowercase__						=				pipeline(UpperCamelCase_   ,						UpperCamelCase_   ,						UpperCamelCase_   ,						jit=UpperCamelCase_							).images
									assert images_eff.shape == (num_samples, 1, 512, 512, 3)
									lowercase__						=				images[2, 0, 256, 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






import json
import logging
import os
import socket

import git
import numpy as np
import torch


logging.basicConfig(
    format='%(asctime)s - %(levelname)s - %(name)s - PID: %(process)d -  %(message)s',
    datefmt='%m/%d/%Y %H:%M:%S',
    level=logging.INFO,
)
lowerCAmelCase	=							logging.getLogger(__name__)
def       _a			(				SCREAMING_SNAKE_CASE					):


		"""simple docstring"""




		lowercase__						=				git.Repo(search_parent_directories=SCREAMING_SNAKE_CASE					)
		lowercase__						=				{
		    '''repo_id''': str(SCREAMING_SNAKE_CASE					),
		    '''repo_sha''': str(repo.head.object.hexsha					),
		    '''repo_branch''': str(repo.active_branch					),
		}

		with open(os.path.join(SCREAMING_SNAKE_CASE						,    '''git_log.json'''					)						,    '''w'''					) as f:
				json.dump(SCREAMING_SNAKE_CASE						,    SCREAMING_SNAKE_CASE						,    indent=4					)
def       _a			(				SCREAMING_SNAKE_CASE					):


		"""simple docstring"""




		if params.n_gpu <= 0:
				lowercase__						=				0
				lowercase__						=				-1
				lowercase__						=				True
				lowercase__						=				False
				return

		assert torch.cuda.is_available()

		logger.info('''Initializing GPUs'''					)
		if params.n_gpu > 1:
				assert params.local_rank != -1

				lowercase__						=				int(os.environ['''WORLD_SIZE''']					)
				lowercase__						=				int(os.environ['''N_GPU_NODE''']					)
				lowercase__						=				int(os.environ['''RANK''']					)

				# number of nodes / node ID
				lowercase__						=				params.world_size // params.n_gpu_per_node
				lowercase__						=				params.global_rank // params.n_gpu_per_node
				lowercase__						=				True

				assert params.n_nodes == int(os.environ['''N_NODES''']					)
				assert params.node_id == int(os.environ['''NODE_RANK''']					)

		# local job (single GPU)
		else:
				assert params.local_rank == -1

				lowercase__						=				1
				lowercase__						=				0
				lowercase__						=				0
				lowercase__						=				0
				lowercase__						=				1
				lowercase__						=				1
				lowercase__						=				False

		# sanity checks
		assert params.n_nodes >= 1
		assert 0 <= params.node_id < params.n_nodes
		assert 0 <= params.local_rank <= params.global_rank < params.world_size
		assert params.world_size == params.n_nodes * params.n_gpu_per_node

		# define whether this is the master process / if we are in multi-node distributed mode
		lowercase__						=				params.node_id == 0 and params.local_rank == 0
		lowercase__						=				params.n_nodes > 1

		# summary
		lowercase__						=				f'--- Global rank: {params.global_rank} - '
		logger.info(PREFIX + '''Number of nodes: %i''' % params.n_nodes					)
		logger.info(PREFIX + '''Node ID        : %i''' % params.node_id					)
		logger.info(PREFIX + '''Local rank     : %i''' % params.local_rank					)
		logger.info(PREFIX + '''World size     : %i''' % params.world_size					)
		logger.info(PREFIX + '''GPUs per node  : %i''' % params.n_gpu_per_node					)
		logger.info(PREFIX + '''Master         : %s''' % str(params.is_master					)					)
		logger.info(PREFIX + '''Multi-node     : %s''' % str(params.multi_node					)					)
		logger.info(PREFIX + '''Multi-GPU      : %s''' % str(params.multi_gpu					)					)
		logger.info(PREFIX + '''Hostname       : %s''' % socket.gethostname()					)

		# set GPU device
		torch.cuda.set_device(params.local_rank					)

		# initialize multi-GPU
		if params.multi_gpu:
				logger.info('''Initializing PyTorch distributed'''					)
				torch.distributed.init_process_group(
				    init_method='''env://'''						,    backend='''nccl'''						,    )
def       _a			(				SCREAMING_SNAKE_CASE					):


		"""simple docstring"""




		np.random.seed(args.seed					)
		torch.manual_seed(args.seed					)
		if args.n_gpu > 0:
				torch.cuda.manual_seed_all(args.seed					)


