



'''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 		_lowercase	(					self :      int  )   ->   Tuple:
    # clean up the VRAM after each test
    super().tearDown()
    gc.collect()




  def 		_lowercase	(					self :      List[str]  )   ->   int:
    __lowerCamelCase			,__lowerCamelCase			:			List[str]       			=					FlaxStableDiffusionPipeline.from_pretrained(
        'stabilityai/stable-diffusion-2'     ,      revision='bf16'     ,      dtype=jnp.bfloataa     ,      )

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

    __lowerCamelCase			:			Optional[Any]       			=					replicate(_a  )
    __lowerCamelCase			:			int       			=					shard(_a  )

    __lowerCamelCase			:			List[str]       			=					jax.random.PRNGKey(0  )
    __lowerCamelCase			:			str       			=					jax.random.split(_a     ,      jax.device_count()  )

    __lowerCamelCase			:			Optional[int]       			=					sd_pipe(_a     ,      _a     ,      _a     ,      num_inference_steps=25     ,      jit=_a  )[0]
    assert images.shape == (jax.device_count(), 1, 768, 768, 3)

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

    __lowerCamelCase			:			int       			=					jnp.asarray(jax.device_get(image_slice.flatten()  )  )
    __lowerCamelCase			:			Optional[int]       			=					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 		_lowercase	(					self :      int  )   ->   Any:
    __lowerCamelCase			:			Tuple       			=					'stabilityai/stable-diffusion-2'
    __lowerCamelCase			,__lowerCamelCase			:			List[Any]       			=					FlaxDPMSolverMultistepScheduler.from_pretrained(_a     ,      subfolder='scheduler'  )
    __lowerCamelCase			,__lowerCamelCase			:			Optional[int]       			=					FlaxStableDiffusionPipeline.from_pretrained(
        _a     ,      scheduler=_a     ,      revision='bf16'     ,      dtype=jnp.bfloataa     ,      )
    __lowerCamelCase			:			Tuple       			=					scheduler_params

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

    __lowerCamelCase			:			int       			=					replicate(_a  )
    __lowerCamelCase			:			List[Any]       			=					shard(_a  )

    __lowerCamelCase			:			List[str]       			=					jax.random.PRNGKey(0  )
    __lowerCamelCase			:			Tuple       			=					jax.random.split(_a     ,      jax.device_count()  )

    __lowerCamelCase			:			Union[str, Any]       			=					sd_pipe(_a     ,      _a     ,      _a     ,      num_inference_steps=25     ,      jit=_a  )[0]
    assert images.shape == (jax.device_count(), 1, 768, 768, 3)

    __lowerCamelCase			:			Optional[int]       			=					images.reshape((images.shape[0] * images.shape[1],) + images.shape[-3:]  )
    __lowerCamelCase			:			Dict       			=					images[0, 253:256, 253:256, -1]

    __lowerCamelCase			:			List[Any]       			=					jnp.asarray(jax.device_get(image_slice.flatten()  )  )
    __lowerCamelCase			:			Optional[int]       			=					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'''





from __future__ import annotations

_UpperCamelCase											=		10

def 			a_	(				_lowerCAmelCase			) -> list[int]:
  __lowerCamelCase			:			str       			=					1
  __lowerCamelCase			:			Union[str, Any]       			=					max(_lowerCAmelCase			)
  while placement <= max_digit:
    # declare and initialize empty buckets
    __lowerCamelCase			:			list[list]       			=					[[] for _ in range(_lowerCAmelCase			)]
    # split list_of_ints between the buckets
    for i in list_of_ints:
      __lowerCamelCase			:			List[str]       			=					int((i / placement) % RADIX			)
      buckets[tmp].append(_lowerCAmelCase			)
    # put each buckets' contents into list_of_ints
    __lowerCamelCase			:			Dict       			=					0
    for b in range(_lowerCAmelCase			):
      for i in buckets[b]:
        __lowerCamelCase			:			List[str]       			=					i
        a += 1
        # move to next
    placement *= RADIX
  return list_of_ints


if __name__ == "__main__":
   import doctest

   doctest.testmod()






