

import gc
import unittest

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


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






@slow
@require_flax
class _UpperCamelCase							(							unittest.TestCase      ):
			'''simple docstring'''







			def 							__UpperCamelCase       (	self							:       Tuple		)							->				int:

				"""simple docstring"""

				super().tearDown()
				gc.collect()
			def 							__UpperCamelCase       (	self							:       str		)							->				int:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE   ,SCREAMING_SNAKE_CASE      :					str    =				FlaxStableDiffusionPipeline.from_pretrained(
				    "stabilityai/stable-diffusion-2"					,					revision="bf16"					,					dtype=jnp.bfloataa					,					)

				SCREAMING_SNAKE_CASE      :					int    =				"A painting of a squirrel eating a burger"
				SCREAMING_SNAKE_CASE      :					Union[str, Any]    =				jax.device_count()
				SCREAMING_SNAKE_CASE      :					Optional[int]    =				num_samples * [prompt]
				SCREAMING_SNAKE_CASE      :					Any    =				sd_pipe.prepare_inputs(a		)

				SCREAMING_SNAKE_CASE      :					Tuple    =				replicate(a		)
				SCREAMING_SNAKE_CASE      :					Dict    =				shard(a		)

				SCREAMING_SNAKE_CASE      :					Union[str, Any]    =				jax.random.PRNGKey(0		)
				SCREAMING_SNAKE_CASE      :					Dict    =				jax.random.split(a					,					jax.device_count()		)

				SCREAMING_SNAKE_CASE      :					str    =				sd_pipe(a					,					a					,					a					,					num_inference_steps=25					,					jit=a		)[0]
				assert images.shape == (jax.device_count(), 1, 768, 768, 3)

				SCREAMING_SNAKE_CASE      :					Tuple    =				images.reshape((images.shape[0] * images.shape[1],) + images.shape[-3:]		)
				SCREAMING_SNAKE_CASE      :					List[str]    =				images[0, 253:256, 253:256, -1]

				SCREAMING_SNAKE_CASE      :					Optional[int]    =				jnp.asarray(jax.device_get(image_slice.flatten()		)		)
				SCREAMING_SNAKE_CASE      :					Any    =				jnp.array([0.4238, 0.4414, 0.4395, 0.4453, 0.4629, 0.4590, 0.4531, 0.4_5508, 0.4512]		)
				print(F"output_slice: {output_slice}"		)
				assert jnp.abs(output_slice - expected_slice		).max() < 1e-2






			def 							__UpperCamelCase       (	self							:       List[str]		)							->				Optional[Any]:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE      :					int    =				"stabilityai/stable-diffusion-2"
				SCREAMING_SNAKE_CASE   ,SCREAMING_SNAKE_CASE      :					Tuple    =				FlaxDPMSolverMultistepScheduler.from_pretrained(a					,					subfolder="scheduler"		)
				SCREAMING_SNAKE_CASE   ,SCREAMING_SNAKE_CASE      :					str    =				FlaxStableDiffusionPipeline.from_pretrained(
				    a					,					scheduler=a					,					revision="bf16"					,					dtype=jnp.bfloataa					,					)
				SCREAMING_SNAKE_CASE      :					Tuple    =				scheduler_params

				SCREAMING_SNAKE_CASE      :					List[str]    =				"A painting of a squirrel eating a burger"
				SCREAMING_SNAKE_CASE      :					Optional[Any]    =				jax.device_count()
				SCREAMING_SNAKE_CASE      :					int    =				num_samples * [prompt]
				SCREAMING_SNAKE_CASE      :					Union[str, Any]    =				sd_pipe.prepare_inputs(a		)

				SCREAMING_SNAKE_CASE      :					Tuple    =				replicate(a		)
				SCREAMING_SNAKE_CASE      :					Optional[int]    =				shard(a		)

				SCREAMING_SNAKE_CASE      :					Dict    =				jax.random.PRNGKey(0		)
				SCREAMING_SNAKE_CASE      :					str    =				jax.random.split(a					,					jax.device_count()		)

				SCREAMING_SNAKE_CASE      :					str    =				sd_pipe(a					,					a					,					a					,					num_inference_steps=25					,					jit=a		)[0]
				assert images.shape == (jax.device_count(), 1, 768, 768, 3)

				SCREAMING_SNAKE_CASE      :					Any    =				images.reshape((images.shape[0] * images.shape[1],) + images.shape[-3:]		)
				SCREAMING_SNAKE_CASE      :					Union[str, Any]    =				images[0, 253:256, 253:256, -1]

				SCREAMING_SNAKE_CASE      :					Union[str, Any]    =				jnp.asarray(jax.device_get(image_slice.flatten()		)		)
				SCREAMING_SNAKE_CASE      :					Union[str, Any]    =				jnp.array([0.4336, 0.4_2969, 0.4453, 0.4199, 0.4297, 0.4531, 0.4434, 0.4434, 0.4297]		)
				print(F"output_slice: {output_slice}"		)
				assert jnp.abs(output_slice - expected_slice		).max() < 1e-2


'''simple docstring'''



import unittest

import numpy as np
import torch

from diffusers import PNDMPipeline, PNDMScheduler, UNetaDModel
from diffusers.utils.testing_utils import enable_full_determinism, require_torch, slow, torch_device


enable_full_determinism()







class 						lowerCAmelCase__   (      unittest.TestCase					):



 """simple docstring"""



 @property
 def       UpperCAmelCase__		(		self					:       List[str]			)  -> Optional[int]:


      """simple docstring"""




      torch.manual_seed(0			)
      __SCREAMING_SNAKE_CASE     =      UNetaDModel(
          block_out_channels=(32, 64)					,							layers_per_block=2					,							sample_size=32					,							in_channels=3					,							out_channels=3					,							down_block_types=("""DownBlock2D""", """AttnDownBlock2D""")					,							up_block_types=("""AttnUpBlock2D""", """UpBlock2D""")					,							)
      return model


 def       UpperCAmelCase__		(		self					:       List[Any]			)  -> List[str]:


      """simple docstring"""




      __SCREAMING_SNAKE_CASE     =      self.dummy_uncond_unet
      __SCREAMING_SNAKE_CASE     =      PNDMScheduler()

      __SCREAMING_SNAKE_CASE     =      PNDMPipeline(unet=__SCREAMING_SNAKE_CASE					,							scheduler=__SCREAMING_SNAKE_CASE			)
      pndm.to(__SCREAMING_SNAKE_CASE			)
      pndm.set_progress_bar_config(disable=__SCREAMING_SNAKE_CASE			)

      __SCREAMING_SNAKE_CASE     =      torch.manual_seed(0			)
      __SCREAMING_SNAKE_CASE     =      pndm(generator=__SCREAMING_SNAKE_CASE					,							num_inference_steps=20					,							output_type="""numpy"""			).images

      __SCREAMING_SNAKE_CASE     =      torch.manual_seed(0			)
      __SCREAMING_SNAKE_CASE     =      pndm(generator=__SCREAMING_SNAKE_CASE					,							num_inference_steps=20					,							output_type="""numpy"""					,							return_dict=__SCREAMING_SNAKE_CASE			)[0]

      __SCREAMING_SNAKE_CASE     =      image[0, -3:, -3:, -1]
      __SCREAMING_SNAKE_CASE     =      image_from_tuple[0, -3:, -3:, -1]

      assert image.shape == (1, 32, 32, 3)
      __SCREAMING_SNAKE_CASE     =      np.array([1.0, 1.0, 0.0, 1.0, 0.0, 1.0, 0.0, 0.0, 0.0]			)

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







@slow
@require_torch
class 						lowerCAmelCase__   (      unittest.TestCase					):



 """simple docstring"""



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


      """simple docstring"""




      __SCREAMING_SNAKE_CASE     =      """google/ddpm-cifar10-32"""

      __SCREAMING_SNAKE_CASE     =      UNetaDModel.from_pretrained(__SCREAMING_SNAKE_CASE			)
      __SCREAMING_SNAKE_CASE     =      PNDMScheduler()

      __SCREAMING_SNAKE_CASE     =      PNDMPipeline(unet=__SCREAMING_SNAKE_CASE					,							scheduler=__SCREAMING_SNAKE_CASE			)
      pndm.to(__SCREAMING_SNAKE_CASE			)
      pndm.set_progress_bar_config(disable=__SCREAMING_SNAKE_CASE			)
      __SCREAMING_SNAKE_CASE     =      torch.manual_seed(0			)
      __SCREAMING_SNAKE_CASE     =      pndm(generator=__SCREAMING_SNAKE_CASE					,							output_type="""numpy"""			).images

      __SCREAMING_SNAKE_CASE     =      image[0, -3:, -3:, -1]

      assert image.shape == (1, 32, 32, 3)
      __SCREAMING_SNAKE_CASE     =      np.array([0.1564, 0.14645, 0.1406, 0.14715, 0.12425, 0.14045, 0.13115, 0.12175, 0.125]			)

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






