

'''simple docstring'''



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



 """simple docstring"""



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


      """simple docstring"""




      super().tearDown()
      gc.collect()
 def  UpperCAmelCase__						(						self   :			Optional[int]     )    ->					Optional[int]:


      """simple docstring"""




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

      __SCREAMING_SNAKE_CASE      = """A painting of a squirrel eating a burger"""
      __SCREAMING_SNAKE_CASE      = jax.device_count()
      __SCREAMING_SNAKE_CASE      = num_samples * [prompt]
      __SCREAMING_SNAKE_CASE      = sd_pipe.prepare_inputs(__SCREAMING_SNAKE_CASE     )

      __SCREAMING_SNAKE_CASE      = replicate(__SCREAMING_SNAKE_CASE     )
      __SCREAMING_SNAKE_CASE      = shard(__SCREAMING_SNAKE_CASE     )

      __SCREAMING_SNAKE_CASE      = jax.random.PRNGKey(0     )
      __SCREAMING_SNAKE_CASE      = jax.random.split(__SCREAMING_SNAKE_CASE   ,				jax.device_count()     )

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

      __SCREAMING_SNAKE_CASE      = images.reshape((images.shape[0] * images.shape[1],) + images.shape[-3:]     )
      __SCREAMING_SNAKE_CASE      = images[0, 253:256, 253:256, -1]

      __SCREAMING_SNAKE_CASE      = jnp.asarray(jax.device_get(image_slice.flatten()     )     )
      __SCREAMING_SNAKE_CASE      = jnp.array([0.4238, 0.4414, 0.4395, 0.4453, 0.4629, 0.4590, 0.4531, 0.45508, 0.4512]     )
      print(f'output_slice: {output_slice}'     )
      assert jnp.abs(output_slice - expected_slice     ).max() < 1E-2


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


      """simple docstring"""




      __SCREAMING_SNAKE_CASE      = """stabilityai/stable-diffusion-2"""
      __SCREAMING_SNAKE_CASE  ,     __SCREAMING_SNAKE_CASE      = FlaxDPMSolverMultistepScheduler.from_pretrained(__SCREAMING_SNAKE_CASE   ,				subfolder="""scheduler"""     )
      __SCREAMING_SNAKE_CASE  ,     __SCREAMING_SNAKE_CASE      = FlaxStableDiffusionPipeline.from_pretrained(
          __SCREAMING_SNAKE_CASE   ,				scheduler=__SCREAMING_SNAKE_CASE   ,				revision="""bf16"""   ,				dtype=jnp.bfloataa   ,				)
      __SCREAMING_SNAKE_CASE      = scheduler_params

      __SCREAMING_SNAKE_CASE      = """A painting of a squirrel eating a burger"""
      __SCREAMING_SNAKE_CASE      = jax.device_count()
      __SCREAMING_SNAKE_CASE      = num_samples * [prompt]
      __SCREAMING_SNAKE_CASE      = sd_pipe.prepare_inputs(__SCREAMING_SNAKE_CASE     )

      __SCREAMING_SNAKE_CASE      = replicate(__SCREAMING_SNAKE_CASE     )
      __SCREAMING_SNAKE_CASE      = shard(__SCREAMING_SNAKE_CASE     )

      __SCREAMING_SNAKE_CASE      = jax.random.PRNGKey(0     )
      __SCREAMING_SNAKE_CASE      = jax.random.split(__SCREAMING_SNAKE_CASE   ,				jax.device_count()     )

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

      __SCREAMING_SNAKE_CASE      = images.reshape((images.shape[0] * images.shape[1],) + images.shape[-3:]     )
      __SCREAMING_SNAKE_CASE      = images[0, 253:256, 253:256, -1]

      __SCREAMING_SNAKE_CASE      = jnp.asarray(jax.device_get(image_slice.flatten()     )     )
      __SCREAMING_SNAKE_CASE      = jnp.array([0.4336, 0.42969, 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 KarrasVePipeline, KarrasVeScheduler, 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   :			int     )    ->					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   :			Optional[Any]     )    ->					Tuple:


      """simple docstring"""




      __SCREAMING_SNAKE_CASE      = self.dummy_uncond_unet
      __SCREAMING_SNAKE_CASE      = KarrasVeScheduler()

      __SCREAMING_SNAKE_CASE      = KarrasVePipeline(unet=__SCREAMING_SNAKE_CASE   ,				scheduler=__SCREAMING_SNAKE_CASE     )
      pipe.to(__SCREAMING_SNAKE_CASE     )
      pipe.set_progress_bar_config(disable=__SCREAMING_SNAKE_CASE     )

      __SCREAMING_SNAKE_CASE      = torch.manual_seed(0     )
      __SCREAMING_SNAKE_CASE      = pipe(num_inference_steps=2   ,				generator=__SCREAMING_SNAKE_CASE   ,				output_type="""numpy"""     ).images

      __SCREAMING_SNAKE_CASE      = torch.manual_seed(0     )
      __SCREAMING_SNAKE_CASE      = pipe(num_inference_steps=2   ,				generator=__SCREAMING_SNAKE_CASE   ,				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([0.0, 1.0, 0.0, 0.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   :			List[Any]     )    ->					Union[str, Any]:


      """simple docstring"""




      __SCREAMING_SNAKE_CASE      = """google/ncsnpp-celebahq-256"""
      __SCREAMING_SNAKE_CASE      = UNetaDModel.from_pretrained(__SCREAMING_SNAKE_CASE     )
      __SCREAMING_SNAKE_CASE      = KarrasVeScheduler()

      __SCREAMING_SNAKE_CASE      = KarrasVePipeline(unet=__SCREAMING_SNAKE_CASE   ,				scheduler=__SCREAMING_SNAKE_CASE     )
      pipe.to(__SCREAMING_SNAKE_CASE     )
      pipe.set_progress_bar_config(disable=__SCREAMING_SNAKE_CASE     )

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

      __SCREAMING_SNAKE_CASE      = image[0, -3:, -3:, -1]
      assert image.shape == (1, 256, 256, 3)
      __SCREAMING_SNAKE_CASE      = np.array([0.578, 0.5811, 0.5924, 0.5809, 0.587, 0.5886, 0.5861, 0.5802, 0.586]     )

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






