





'''simple docstring'''

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







  def 				UpperCAmelCase_				(		self		)		->						Optional[Any]:






        """simple docstring"""





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

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







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






        """simple docstring"""





        _snake_case     ,							_snake_case		:	Tuple					     =     FlaxStableDiffusionPipeline.from_pretrained(
            '''hf-internal-testing/tiny-stable-diffusion-pipe'''     ,     safety_checker=__lowercase		)

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

        _snake_case		:	Union[str, Any]					     =     jax.random.PRNGKey(0		)
        _snake_case		:	Any					     =     4

        _snake_case		:	Tuple					     =     jax.device_count()
        _snake_case		:	List[str]					     =     num_samples * [prompt]
        _snake_case		:	List[str]					     =     pipeline.prepare_inputs(__lowercase		)

        # shard inputs and rng
        _snake_case		:	Optional[int]					     =     replicate(__lowercase		)
        _snake_case		:	Tuple					     =     jax.random.split(__lowercase     ,     __lowercase		)
        _snake_case		:	Optional[int]					     =     shard(__lowercase		)

        _snake_case		:	Optional[int]					     =     pipeline(__lowercase     ,     __lowercase     ,     __lowercase     ,     __lowercase     ,     jit=__lowercase		).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_514_745		) < 1E-3
              assert np.abs(np.abs(__lowercase     ,     dtype=np.floataa		).sum() - 49_947.875		) < 5E-1

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







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






        """simple docstring"""





        _snake_case     ,							_snake_case		:	Dict					     =     FlaxStableDiffusionPipeline.from_pretrained(
            '''CompVis/stable-diffusion-v1-4'''     ,     revision='''flax'''     ,     safety_checker=__lowercase		)

        _snake_case		:	List[str]					     =     (
            '''A cinematic film still of Morgan Freeman starring as Jimi Hendrix, portrait, 40mm lens, shallow depth of'''
            ''' field, close up, split lighting, cinematic'''
        )

        _snake_case		:	Dict					     =     jax.random.PRNGKey(0		)
        _snake_case		:	Optional[int]					     =     50

        _snake_case		:	str					     =     jax.device_count()
        _snake_case		:	str					     =     num_samples * [prompt]
        _snake_case		:	List[str]					     =     pipeline.prepare_inputs(__lowercase		)

        # shard inputs and rng
        _snake_case		:	Optional[int]					     =     replicate(__lowercase		)
        _snake_case		:	List[Any]					     =     jax.random.split(__lowercase     ,     __lowercase		)
        _snake_case		:	Any					     =     shard(__lowercase		)

        _snake_case		:	Tuple					     =     pipeline(__lowercase     ,     __lowercase     ,     __lowercase     ,     __lowercase     ,     jit=__lowercase		).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_652_401)		) < 1E-3
              assert np.abs((np.abs(__lowercase     ,     dtype=np.floataa		).sum() - 2_383_808.2)		) < 5E-1







  def 				UpperCAmelCase_				(		self		)		->						Tuple:






        """simple docstring"""





        _snake_case     ,							_snake_case		:	List[Any]					     =     FlaxStableDiffusionPipeline.from_pretrained(
            '''CompVis/stable-diffusion-v1-4'''     ,     revision='''bf16'''     ,     dtype=jnp.bfloataa     ,     safety_checker=__lowercase		)

        _snake_case		:	List[str]					     =     (
            '''A cinematic film still of Morgan Freeman starring as Jimi Hendrix, portrait, 40mm lens, shallow depth of'''
            ''' field, close up, split lighting, cinematic'''
        )

        _snake_case		:	Tuple					     =     jax.random.PRNGKey(0		)
        _snake_case		:	Tuple					     =     50

        _snake_case		:	Union[str, Any]					     =     jax.device_count()
        _snake_case		:	Dict					     =     num_samples * [prompt]
        _snake_case		:	List[Any]					     =     pipeline.prepare_inputs(__lowercase		)

        # shard inputs and rng
        _snake_case		:	Any					     =     replicate(__lowercase		)
        _snake_case		:	Dict					     =     jax.random.split(__lowercase     ,     __lowercase		)
        _snake_case		:	Dict					     =     shard(__lowercase		)

        _snake_case		:	int					     =     pipeline(__lowercase     ,     __lowercase     ,     __lowercase     ,     __lowercase     ,     jit=__lowercase		).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_003_906)		) < 1E-3
              assert np.abs((np.abs(__lowercase     ,     dtype=np.floataa		).sum() - 2_373_516.75)		) < 5E-1







  def 				UpperCAmelCase_				(		self		)		->						str:






        """simple docstring"""





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

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

        _snake_case		:	Union[str, Any]					     =     jax.random.PRNGKey(0		)
        _snake_case		:	Dict					     =     50

        _snake_case		:	List[Any]					     =     jax.device_count()
        _snake_case		:	Optional[int]					     =     num_samples * [prompt]
        _snake_case		:	Dict					     =     pipeline.prepare_inputs(__lowercase		)

        # shard inputs and rng
        _snake_case		:	Optional[Any]					     =     replicate(__lowercase		)
        _snake_case		:	Optional[int]					     =     jax.random.split(__lowercase     ,     __lowercase		)
        _snake_case		:	List[Any]					     =     shard(__lowercase		)

        _snake_case		:	Tuple					     =     pipeline(__lowercase     ,     __lowercase     ,     __lowercase     ,     __lowercase     ,     jit=__lowercase		).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_003_906)		) < 1E-3
              assert np.abs((np.abs(__lowercase     ,     dtype=np.floataa		).sum() - 2_373_516.75)		) < 5E-1







  def 				UpperCAmelCase_				(		self		)		->						Dict:






        """simple docstring"""





        _snake_case		:	List[Any]					     =     FlaxDDIMScheduler(
            beta_start=0.00_085     ,     beta_end=0.012     ,     beta_schedule='''scaled_linear'''     ,     set_alpha_to_one=__lowercase     ,     steps_offset=1     ,     )

        _snake_case     ,							_snake_case		:	List[Any]					     =     FlaxStableDiffusionPipeline.from_pretrained(
            '''CompVis/stable-diffusion-v1-4'''     ,     revision='''bf16'''     ,     dtype=jnp.bfloataa     ,     scheduler=__lowercase     ,     safety_checker=__lowercase     ,     )
        _snake_case		:	List[str]					     =     scheduler.create_state()

        _snake_case		:	List[Any]					     =     scheduler_state

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

        _snake_case		:	Any					     =     jax.random.PRNGKey(0		)
        _snake_case		:	List[str]					     =     50

        _snake_case		:	Optional[Any]					     =     jax.device_count()
        _snake_case		:	Optional[Any]					     =     num_samples * [prompt]
        _snake_case		:	List[Any]					     =     pipeline.prepare_inputs(__lowercase		)

        # shard inputs and rng
        _snake_case		:	Union[str, Any]					     =     replicate(__lowercase		)
        _snake_case		:	Optional[int]					     =     jax.random.split(__lowercase     ,     __lowercase		)
        _snake_case		:	Dict					     =     shard(__lowercase		)

        _snake_case		:	List[str]					     =     pipeline(__lowercase     ,     __lowercase     ,     __lowercase     ,     __lowercase     ,     jit=__lowercase		).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.045_043_945)		) < 1E-3
              assert np.abs((np.abs(__lowercase     ,     dtype=np.floataa		).sum() - 2_347_693.5)		) < 5E-1







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






        """simple docstring"""





        _snake_case		:	List[str]					     =     (
            '''A cinematic film still of Morgan Freeman starring as Jimi Hendrix, portrait, 40mm lens, shallow depth of'''
            ''' field, close up, split lighting, cinematic'''
        )

        _snake_case		:	List[str]					     =     jax.device_count()
        _snake_case		:	Optional[Any]					     =     num_samples * [prompt]
        _snake_case		:	str					     =     jax.random.split(jax.random.PRNGKey(0		)     ,     __lowercase		)

        _snake_case     ,							_snake_case		:	Any					     =     FlaxStableDiffusionPipeline.from_pretrained(
            '''CompVis/stable-diffusion-v1-4'''     ,     revision='''bf16'''     ,     dtype=jnp.bfloataa     ,     safety_checker=__lowercase     ,     )

        _snake_case		:	Dict					     =     replicate(__lowercase		)
        _snake_case		:	Optional[Any]					     =     pipeline.prepare_inputs(__lowercase		)
        _snake_case		:	List[Any]					     =     shard(__lowercase		)
        _snake_case		:	int					     =     pipeline(__lowercase     ,     __lowercase     ,     __lowercase     ,     jit=__lowercase		).images
        assert images.shape == (num_samples, 1, 512, 512, 3)
        _snake_case		:	Any					     =     images[2, 0, 256, 10:17, 1]

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

        _snake_case		:	Union[str, Any]					     =     replicate(__lowercase		)
        _snake_case		:	Optional[int]					     =     pipeline.prepare_inputs(__lowercase		)
        _snake_case		:	Union[str, Any]					     =     shard(__lowercase		)
        _snake_case		:	int					     =     pipeline(__lowercase     ,     __lowercase     ,     __lowercase     ,     jit=__lowercase		).images
        assert images_eff.shape == (num_samples, 1, 512, 512, 3)
        _snake_case		:	Any					     =     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









'''simple docstring'''

from typing import List, Optional, Union

import numpy as np
import PIL.Image

from ...image_processing_utils import BaseImageProcessor, BatchFeature
from ...image_transforms import rescale, resize, to_channel_dimension_format
from ...image_utils import (
    ChannelDimension,
    PILImageResampling,
    get_image_size,
    make_list_of_images,
    to_numpy_array,
    valid_images,
)
from ...utils import TensorType, logging


UpperCAmelCase					:		Dict		=							logging.get_logger(__name__)





class lowerCamelCase						(a__					):
  _lowercase  :	int       					=				["""pixel_values"""]







  def __init__(		self     ,     lowercase__ = True     ,     lowercase__ = 32     ,     lowercase__=PILImageResampling.BILINEAR     ,     lowercase__ = True     ,     **lowercase__     ,     )		->						None:






        """simple docstring"""





        _snake_case		:	Any					     =     do_resize
        _snake_case		:	List[str]					     =     do_rescale
        _snake_case		:	Any					     =     size_divisor
        _snake_case		:	Optional[Any]					     =     resample
        super().__init__(**lowercase__		)







  def 				UpperCAmelCase_				(		self     ,     lowercase__     ,     lowercase__     ,     lowercase__     ,     lowercase__ = None     ,     **lowercase__		)		->						np.ndarray:






        """simple docstring"""





        _snake_case     ,							_snake_case		:	Dict					     =     get_image_size(lowercase__		)
        # Rounds the height and width down to the closest multiple of size_divisor
        _snake_case		:	Optional[int]					     =     height // size_divisor * size_divisor
        _snake_case		:	Dict					     =     width // size_divisor * size_divisor
        _snake_case		:	str					     =     resize(lowercase__     ,     (new_h, new_w)     ,     resample=lowercase__     ,     data_format=lowercase__     ,     **lowercase__		)
        return image







  def 				UpperCAmelCase_				(		self     ,     lowercase__     ,     lowercase__     ,     lowercase__ = None     ,     **lowercase__		)		->						np.ndarray:






        """simple docstring"""





        return rescale(image=lowercase__     ,     scale=lowercase__     ,     data_format=lowercase__     ,     **lowercase__		)







  def 				UpperCAmelCase_				(		self     ,     lowercase__     ,     lowercase__ = None     ,     lowercase__ = None     ,     lowercase__=None     ,     lowercase__ = None     ,     lowercase__ = None     ,     lowercase__ = ChannelDimension.FIRST     ,     **lowercase__     ,     )		->						BatchFeature:






        """simple docstring"""





        _snake_case		:	Any					     =     do_resize if do_resize is not None else self.do_resize
        _snake_case		:	List[Any]					     =     do_rescale if do_rescale is not None else self.do_rescale
        _snake_case		:	List[str]					     =     size_divisor if size_divisor is not None else self.size_divisor
        _snake_case		:	int					     =     resample if resample is not None else self.resample

        if do_resize and size_divisor is None:
              raise ValueError('''size_divisor is required for resizing'''		)

        _snake_case		:	Tuple					     =     make_list_of_images(lowercase__		)

        if not valid_images(lowercase__		):
              raise ValueError('''Invalid image(s)'''		)

        # All transformations expect numpy arrays.
        _snake_case		:	Tuple					     =     [to_numpy_array(lowercase__		) for img in images]

        if do_resize:
              _snake_case		:	Optional[int]					     =     [self.resize(lowercase__     ,     size_divisor=lowercase__     ,     resample=lowercase__		) for image in images]

        if do_rescale:
              _snake_case		:	Union[str, Any]					     =     [self.rescale(lowercase__     ,     scale=1 / 255		) for image in images]

        _snake_case		:	Union[str, Any]					     =     [to_channel_dimension_format(lowercase__     ,     lowercase__		) for image in images]

        _snake_case		:	List[str]					     =     {'''pixel_values''': images}
        return BatchFeature(data=lowercase__     ,     tensor_type=lowercase__		)


