






import gc
import math
import unittest

import torch

from diffusers import UNetaDModel
from diffusers.utils import floats_tensor, logging, slow, torch_all_close, torch_device
from diffusers.utils.testing_utils import enable_full_determinism

from .test_modeling_common import ModelTesterMixin, UNetTesterMixin


_a      	=  logging.get_logger(__name__)

enable_full_determinism()





class __lowerCamelCase (      snake_case__     ,       snake_case__     ,       unittest.TestCase):





	"""simple docstring"""






	UpperCamelCase__								=      UNetaDModel
	UpperCamelCase__								=      "sample"

	@property
	def        UpperCamelCase    (							self    ):
							"""simple docstring"""







							_UpperCAmelCase									=			4
							_UpperCAmelCase									=			3
							_UpperCAmelCase									=			(32, 32)

							_UpperCAmelCase									=			floats_tensor((batch_size, num_channels) + sizes    ).to(UpperCAmelCase    )
							_UpperCAmelCase									=			torch.tensor([10]    ).to(UpperCAmelCase    )

							return {"sample": noise, "timestep": time_step}

	@property
	def        UpperCamelCase    (							self    ):
							"""simple docstring"""







							return (3, 32, 32)

	@property
	def        UpperCamelCase    (							self    ):
							"""simple docstring"""







							return (3, 32, 32)

	def        UpperCamelCase    (							self    ):
							"""simple docstring"""







							_UpperCAmelCase									=			{
							    'block_out_channels': (32, 64),
							    'down_block_types': ('DownBlock2D', 'AttnDownBlock2D'),
							    'up_block_types': ('AttnUpBlock2D', 'UpBlock2D'),
							    'attention_head_dim': 3,
							    'out_channels': 3,
							    'in_channels': 3,
							    'layers_per_block': 2,
							    'sample_size': 32,
							}
							_UpperCAmelCase									=			self.dummy_input
							return init_dict, inputs_dict





class __lowerCamelCase (      snake_case__     ,       snake_case__     ,       unittest.TestCase):





	"""simple docstring"""






	UpperCamelCase__								=      UNetaDModel
	UpperCamelCase__								=      "sample"

	@property
	def        UpperCamelCase    (							self    ):
							"""simple docstring"""







							_UpperCAmelCase									=			4
							_UpperCAmelCase									=			4
							_UpperCAmelCase									=			(32, 32)

							_UpperCAmelCase									=			floats_tensor((batch_size, num_channels) + sizes    ).to(UpperCAmelCase    )
							_UpperCAmelCase									=			torch.tensor([10]    ).to(UpperCAmelCase    )

							return {"sample": noise, "timestep": time_step}

	@property
	def        UpperCamelCase    (							self    ):
							"""simple docstring"""







							return (4, 32, 32)

	@property
	def        UpperCamelCase    (							self    ):
							"""simple docstring"""







							return (4, 32, 32)

	def        UpperCamelCase    (							self    ):
							"""simple docstring"""







							_UpperCAmelCase									=			{
							    'sample_size': 32,
							    'in_channels': 4,
							    'out_channels': 4,
							    'layers_per_block': 2,
							    'block_out_channels': (32, 64),
							    'attention_head_dim': 32,
							    'down_block_types': ('DownBlock2D', 'DownBlock2D'),
							    'up_block_types': ('UpBlock2D', 'UpBlock2D'),
							}
							_UpperCAmelCase									=			self.dummy_input
							return init_dict, inputs_dict

	def        UpperCamelCase    (							self    ):
							"""simple docstring"""







							_UpperCAmelCase					, _UpperCAmelCase									=			UNetaDModel.from_pretrained('fusing/unet-ldm-dummy-update'       ,    output_loading_info=UpperCAmelCase    )

							self.assertIsNotNone(UpperCAmelCase    )
							self.assertEqual(len(loading_info['missing_keys']    )       ,    0    )

							model.to(UpperCAmelCase    )
							_UpperCAmelCase									=			model(**self.dummy_input    ).sample

							assert image is not None, "Make sure output is not None"

	@unittest.skipIf(torch_device != 'cuda'       ,    'This test is supposed to run on GPU'    )
	def        UpperCamelCase    (							self    ):
							"""simple docstring"""







							_UpperCAmelCase					, _UpperCAmelCase									=			UNetaDModel.from_pretrained('fusing/unet-ldm-dummy-update'       ,    output_loading_info=UpperCAmelCase    )
							model.to(UpperCAmelCase    )
							_UpperCAmelCase									=			model(**self.dummy_input    ).sample

							assert image is not None, "Make sure output is not None"

	@unittest.skipIf(torch_device != 'cuda'       ,    'This test is supposed to run on GPU'    )
	def        UpperCamelCase    (							self    ):
							"""simple docstring"""







							_UpperCAmelCase					, _UpperCAmelCase									=			UNetaDModel.from_pretrained('fusing/unet-ldm-dummy-update'       ,    output_loading_info=UpperCAmelCase    )
							model_accelerate.to(UpperCAmelCase    )
							model_accelerate.eval()

							_UpperCAmelCase									=			torch.randn(
							    1       ,    model_accelerate.config.in_channels       ,    model_accelerate.config.sample_size       ,    model_accelerate.config.sample_size       ,    generator=torch.manual_seed(0    )       ,    )
							_UpperCAmelCase									=			noise.to(UpperCAmelCase    )
							_UpperCAmelCase									=			torch.tensor([10] * noise.shape[0]    ).to(UpperCAmelCase    )

							_UpperCAmelCase									=			model_accelerate(UpperCAmelCase       ,    UpperCAmelCase    )['sample']

							# two models don't need to stay in the device at the same time
							del model_accelerate
							torch.cuda.empty_cache()
							gc.collect()

							_UpperCAmelCase					, _UpperCAmelCase									=			UNetaDModel.from_pretrained(
							    'fusing/unet-ldm-dummy-update'       ,    output_loading_info=UpperCAmelCase       ,    low_cpu_mem_usage=UpperCAmelCase    )
							model_normal_load.to(UpperCAmelCase    )
							model_normal_load.eval()
							_UpperCAmelCase									=			model_normal_load(UpperCAmelCase       ,    UpperCAmelCase    )['sample']

							assert torch_all_close(UpperCAmelCase       ,    UpperCAmelCase       ,    rtol=1e-3    )

	def        UpperCamelCase    (							self    ):
							"""simple docstring"""







							_UpperCAmelCase									=			UNetaDModel.from_pretrained('fusing/unet-ldm-dummy-update'    )
							model.eval()
							model.to(UpperCAmelCase    )

							_UpperCAmelCase									=			torch.randn(
							    1       ,    model.config.in_channels       ,    model.config.sample_size       ,    model.config.sample_size       ,    generator=torch.manual_seed(0    )       ,    )
							_UpperCAmelCase									=			noise.to(UpperCAmelCase    )
							_UpperCAmelCase									=			torch.tensor([10] * noise.shape[0]    ).to(UpperCAmelCase    )

							with torch.no_grad():
													_UpperCAmelCase									=			model(UpperCAmelCase       ,    UpperCAmelCase    ).sample

							_UpperCAmelCase									=			output[0, -1, -3:, -3:].flatten().cpu()
							# fmt: off
							_UpperCAmelCase									=			torch.tensor([-13.32_58, -20.11_00, -15.98_73, -17.66_17, -23.05_96, -17.94_19, -13.36_75, -16.18_89, -12.38_00]    )
							# fmt: on

							self.assertTrue(torch_all_close(UpperCAmelCase       ,    UpperCAmelCase       ,    rtol=1e-3    )    )





class __lowerCamelCase (      snake_case__     ,       snake_case__     ,       unittest.TestCase):





	"""simple docstring"""






	UpperCamelCase__								=      UNetaDModel
	UpperCamelCase__								=      "sample"

	@property
	def        UpperCamelCase    (							self       ,    UpperCAmelCase=(32, 32)    ):
							"""simple docstring"""







							_UpperCAmelCase									=			4
							_UpperCAmelCase									=			3

							_UpperCAmelCase									=			floats_tensor((batch_size, num_channels) + sizes    ).to(UpperCAmelCase    )
							_UpperCAmelCase									=			torch.tensor(batch_size * [10]    ).to(dtype=torch.intaa       ,    device=UpperCAmelCase    )

							return {"sample": noise, "timestep": time_step}

	@property
	def        UpperCamelCase    (							self    ):
							"""simple docstring"""







							return (3, 32, 32)

	@property
	def        UpperCamelCase    (							self    ):
							"""simple docstring"""







							return (3, 32, 32)

	def        UpperCamelCase    (							self    ):
							"""simple docstring"""







							_UpperCAmelCase									=			{
							    'block_out_channels': [32, 64, 64, 64],
							    'in_channels': 3,
							    'layers_per_block': 1,
							    'out_channels': 3,
							    'time_embedding_type': 'fourier',
							    'norm_eps': 1e-6,
							    'mid_block_scale_factor': math.sqrt(2.0    ),
							    'norm_num_groups': None,
							    'down_block_types': [
							        'SkipDownBlock2D',
							        'AttnSkipDownBlock2D',
							        'SkipDownBlock2D',
							        'SkipDownBlock2D',
							    ],
							    'up_block_types': [
							        'SkipUpBlock2D',
							        'SkipUpBlock2D',
							        'AttnSkipUpBlock2D',
							        'SkipUpBlock2D',
							    ],
							}
							_UpperCAmelCase									=			self.dummy_input
							return init_dict, inputs_dict

	@slow
	def        UpperCamelCase    (							self    ):
							"""simple docstring"""







							_UpperCAmelCase					, _UpperCAmelCase									=			UNetaDModel.from_pretrained('google/ncsnpp-celebahq-256'       ,    output_loading_info=UpperCAmelCase    )
							self.assertIsNotNone(UpperCAmelCase    )
							self.assertEqual(len(loading_info['missing_keys']    )       ,    0    )

							model.to(UpperCAmelCase    )
							_UpperCAmelCase									=			self.dummy_input
							_UpperCAmelCase									=			floats_tensor((4, 3) + (256, 256)    ).to(UpperCAmelCase    )
							_UpperCAmelCase									=			noise
							_UpperCAmelCase									=			model(**UpperCAmelCase    )

							assert image is not None, "Make sure output is not None"

	@slow
	def        UpperCamelCase    (							self    ):
							"""simple docstring"""







							_UpperCAmelCase									=			UNetaDModel.from_pretrained('google/ncsnpp-celebahq-256'    )
							model.to(UpperCAmelCase    )

							_UpperCAmelCase									=			4
							_UpperCAmelCase									=			3
							_UpperCAmelCase									=			(256, 256)

							_UpperCAmelCase									=			torch.ones((batch_size, num_channels) + sizes    ).to(UpperCAmelCase    )
							_UpperCAmelCase									=			torch.tensor(batch_size * [1e-4]    ).to(UpperCAmelCase    )

							with torch.no_grad():
													_UpperCAmelCase									=			model(UpperCAmelCase       ,    UpperCAmelCase    ).sample

							_UpperCAmelCase									=			output[0, -3:, -3:, -1].flatten().cpu()
							# fmt: off
							_UpperCAmelCase									=			torch.tensor([-48_42.86_91, -64_99.66_31, -38_00.19_53, -79_78.26_86, -1_09_80.71_29, -2_00_28.85_35, 81_48.28_22, 23_42.29_05, 5_67.76_08]    )
							# fmt: on

							self.assertTrue(torch_all_close(UpperCAmelCase       ,    UpperCAmelCase       ,    rtol=1e-2    )    )

	def        UpperCamelCase    (							self    ):
							"""simple docstring"""







							_UpperCAmelCase									=			UNetaDModel.from_pretrained('fusing/ncsnpp-ffhq-ve-dummy-update'    )
							model.to(UpperCAmelCase    )

							_UpperCAmelCase									=			4
							_UpperCAmelCase									=			3
							_UpperCAmelCase									=			(32, 32)

							_UpperCAmelCase									=			torch.ones((batch_size, num_channels) + sizes    ).to(UpperCAmelCase    )
							_UpperCAmelCase									=			torch.tensor(batch_size * [1e-4]    ).to(UpperCAmelCase    )

							with torch.no_grad():
													_UpperCAmelCase									=			model(UpperCAmelCase       ,    UpperCAmelCase    ).sample

							_UpperCAmelCase									=			output[0, -3:, -3:, -1].flatten().cpu()
							# fmt: off
							_UpperCAmelCase									=			torch.tensor([-0.03_25, -0.09_00, -0.08_69, -0.03_32, -0.07_25, -0.02_70, -0.01_01, 0.02_27, 0.02_56]    )
							# fmt: on

							self.assertTrue(torch_all_close(UpperCAmelCase       ,    UpperCAmelCase       ,    rtol=1e-2    )    )

	def        UpperCamelCase    (							self    ):
							"""simple docstring"""







							pass








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 __lowerCamelCase (      unittest.TestCase):





	"""simple docstring"""






	def        UpperCamelCase    (							self    ):
							"""simple docstring"""







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

													_UpperCAmelCase									=			[t[-1] for t in os.walk(os.path.join(UpperCAmelCase       ,    os.listdir(UpperCAmelCase    )[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 __lowerCamelCase (      unittest.TestCase):





	"""simple docstring"""






	def        UpperCamelCase    (							self    ):
							"""simple docstring"""







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

							_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(UpperCAmelCase    )

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

							_UpperCAmelCase									=			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.1_51_47_45    ) < 1e-3
													assert np.abs(np.abs(UpperCAmelCase       ,    dtype=np.floataa    ).sum() - 4_99_47.8_75    ) < 5e-1

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

	def        UpperCamelCase    (							self    ):
							"""simple docstring"""







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

							_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(UpperCAmelCase    )

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

							_UpperCAmelCase									=			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.05_65_24_01)    ) < 1e-3
													assert np.abs((np.abs(UpperCAmelCase       ,    dtype=np.floataa    ).sum() - 2_38_38_08.2)    ) < 5e-1

	def        UpperCamelCase    (							self    ):
							"""simple docstring"""







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

							_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(UpperCAmelCase    )

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

							_UpperCAmelCase									=			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.04_00_39_06)    ) < 1e-3
													assert np.abs((np.abs(UpperCAmelCase       ,    dtype=np.floataa    ).sum() - 2_37_35_16.75)    ) < 5e-1

	def        UpperCamelCase    (							self    ):
							"""simple docstring"""







							_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(UpperCAmelCase    )

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

							_UpperCAmelCase									=			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.04_00_39_06)    ) < 1e-3
													assert np.abs((np.abs(UpperCAmelCase       ,    dtype=np.floataa    ).sum() - 2_37_35_16.75)    ) < 5e-1

	def        UpperCamelCase    (							self    ):
							"""simple docstring"""







							_UpperCAmelCase									=			FlaxDDIMScheduler(
							    beta_start=0.0_00_85       ,    beta_end=0.0_12       ,    beta_schedule='scaled_linear'       ,    set_alpha_to_one=UpperCAmelCase       ,    steps_offset=1       ,    )

							_UpperCAmelCase					, _UpperCAmelCase									=			FlaxStableDiffusionPipeline.from_pretrained(
							    'CompVis/stable-diffusion-v1-4'       ,    revision='bf16'       ,    dtype=jnp.bfloataa       ,    scheduler=UpperCAmelCase       ,    safety_checker=UpperCAmelCase       ,    )
							_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(UpperCAmelCase    )

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

							_UpperCAmelCase									=			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.0_45_04_39_45)    ) < 1e-3
													assert np.abs((np.abs(UpperCAmelCase       ,    dtype=np.floataa    ).sum() - 2_34_76_93.5)    ) < 5e-1

	def        UpperCamelCase    (							self    ):
							"""simple docstring"""







							_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    )       ,    UpperCAmelCase    )

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

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

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

							_UpperCAmelCase									=			replicate(UpperCAmelCase    )
							_UpperCAmelCase									=			pipeline.prepare_inputs(UpperCAmelCase    )
							_UpperCAmelCase									=			shard(UpperCAmelCase    )
							_UpperCAmelCase									=			pipeline(UpperCAmelCase       ,    UpperCAmelCase       ,    UpperCAmelCase       ,    jit=UpperCAmelCase    ).images
							assert images_eff.shape == (num_samples, 1, 512, 512, 3)
							_UpperCAmelCase									=			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
