



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




      """simple docstring"""




      def      SCREAMING_SNAKE_CASE__ (      self				:      str       ):
        with tempfile.TemporaryDirectory() as tmpdirname:
          # pipeline has Flax weights
          lowerCAmelCase_   :				int								= FlaxDiffusionPipeline.from_pretrained(
              'hf-internal-testing/tiny-stable-diffusion-pipe'    ,			safety_checker=SCREAMING_SNAKE_CASE_    ,			cache_dir=SCREAMING_SNAKE_CASE_       )

          lowerCAmelCase_   :				Optional[Any]								= [t[-1] for t in os.walk(os.path.join(SCREAMING_SNAKE_CASE_    ,			os.listdir(SCREAMING_SNAKE_CASE_       )[0]    ,			'snapshots'       )       )]
          lowerCAmelCase_   :				str								= [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       UpperCamelCase__   (	unittest.TestCase  ):




      """simple docstring"""




      def      SCREAMING_SNAKE_CASE__ (      self				:      Tuple       ):
        lowerCAmelCase_						,lowerCAmelCase_   :				Optional[int]								= FlaxStableDiffusionPipeline.from_pretrained(
            'hf-internal-testing/tiny-stable-diffusion-pipe'    ,			safety_checker=SCREAMING_SNAKE_CASE_       )

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

        lowerCAmelCase_   :				Any								= jax.random.PRNGKey(0       )
        lowerCAmelCase_   :				Tuple								= 4

        lowerCAmelCase_   :				List[Any]								= jax.device_count()
        lowerCAmelCase_   :				List[Any]								= num_samples * [prompt]
        lowerCAmelCase_   :				Any								= pipeline.prepare_inputs(SCREAMING_SNAKE_CASE_       )

        # shard inputs and rng
        lowerCAmelCase_   :				Optional[Any]								= replicate(SCREAMING_SNAKE_CASE_       )
        lowerCAmelCase_   :				Tuple								= jax.random.split(SCREAMING_SNAKE_CASE_    ,			SCREAMING_SNAKE_CASE_       )
        lowerCAmelCase_   :				Union[str, Any]								= shard(SCREAMING_SNAKE_CASE_       )

        lowerCAmelCase_   :				List[str]								= pipeline(SCREAMING_SNAKE_CASE_    ,			SCREAMING_SNAKE_CASE_    ,			SCREAMING_SNAKE_CASE_    ,			SCREAMING_SNAKE_CASE_    ,			jit=SCREAMING_SNAKE_CASE_       ).images

        assert images.shape == (num_samples, 1, 6_4, 6_4, 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(SCREAMING_SNAKE_CASE_    ,			dtype=np.floataa       ).sum() - 4_99_47.8_75       ) < 5E-1

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



      def      SCREAMING_SNAKE_CASE__ (      self				:      Dict       ):
        lowerCAmelCase_						,lowerCAmelCase_   :				Union[str, Any]								= FlaxStableDiffusionPipeline.from_pretrained(
            'CompVis/stable-diffusion-v1-4'    ,			revision='flax'    ,			safety_checker=SCREAMING_SNAKE_CASE_       )

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

        lowerCAmelCase_   :				List[Any]								= jax.random.PRNGKey(0       )
        lowerCAmelCase_   :				Union[str, Any]								= 5_0

        lowerCAmelCase_   :				str								= jax.device_count()
        lowerCAmelCase_   :				Any								= num_samples * [prompt]
        lowerCAmelCase_   :				Tuple								= pipeline.prepare_inputs(SCREAMING_SNAKE_CASE_       )

        # shard inputs and rng
        lowerCAmelCase_   :				List[str]								= replicate(SCREAMING_SNAKE_CASE_       )
        lowerCAmelCase_   :				Optional[Any]								= jax.random.split(SCREAMING_SNAKE_CASE_    ,			SCREAMING_SNAKE_CASE_       )
        lowerCAmelCase_   :				str								= shard(SCREAMING_SNAKE_CASE_       )

        lowerCAmelCase_   :				List[Any]								= pipeline(SCREAMING_SNAKE_CASE_    ,			SCREAMING_SNAKE_CASE_    ,			SCREAMING_SNAKE_CASE_    ,			SCREAMING_SNAKE_CASE_    ,			jit=SCREAMING_SNAKE_CASE_       ).images

        assert images.shape == (num_samples, 1, 5_1_2, 5_1_2, 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(SCREAMING_SNAKE_CASE_    ,			dtype=np.floataa       ).sum() - 2_38_38_08.2)       ) < 5E-1



      def      SCREAMING_SNAKE_CASE__ (      self				:      int       ):
        lowerCAmelCase_						,lowerCAmelCase_   :				Optional[int]								= FlaxStableDiffusionPipeline.from_pretrained(
            'CompVis/stable-diffusion-v1-4'    ,			revision='bf16'    ,			dtype=jnp.bfloataa    ,			safety_checker=SCREAMING_SNAKE_CASE_       )

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

        lowerCAmelCase_   :				Optional[int]								= jax.random.PRNGKey(0       )
        lowerCAmelCase_   :				List[Any]								= 5_0

        lowerCAmelCase_   :				List[Any]								= jax.device_count()
        lowerCAmelCase_   :				Optional[Any]								= num_samples * [prompt]
        lowerCAmelCase_   :				Union[str, Any]								= pipeline.prepare_inputs(SCREAMING_SNAKE_CASE_       )

        # shard inputs and rng
        lowerCAmelCase_   :				List[Any]								= replicate(SCREAMING_SNAKE_CASE_       )
        lowerCAmelCase_   :				List[Any]								= jax.random.split(SCREAMING_SNAKE_CASE_    ,			SCREAMING_SNAKE_CASE_       )
        lowerCAmelCase_   :				Dict								= shard(SCREAMING_SNAKE_CASE_       )

        lowerCAmelCase_   :				List[Any]								= pipeline(SCREAMING_SNAKE_CASE_    ,			SCREAMING_SNAKE_CASE_    ,			SCREAMING_SNAKE_CASE_    ,			SCREAMING_SNAKE_CASE_    ,			jit=SCREAMING_SNAKE_CASE_       ).images

        assert images.shape == (num_samples, 1, 5_1_2, 5_1_2, 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(SCREAMING_SNAKE_CASE_    ,			dtype=np.floataa       ).sum() - 2_37_35_16.75)       ) < 5E-1



      def      SCREAMING_SNAKE_CASE__ (      self				:      Union[str, Any]       ):
        lowerCAmelCase_						,lowerCAmelCase_   :				str								= FlaxStableDiffusionPipeline.from_pretrained(
            'CompVis/stable-diffusion-v1-4'    ,			revision='bf16'    ,			dtype=jnp.bfloataa       )

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

        lowerCAmelCase_   :				Optional[Any]								= jax.random.PRNGKey(0       )
        lowerCAmelCase_   :				List[str]								= 5_0

        lowerCAmelCase_   :				Tuple								= jax.device_count()
        lowerCAmelCase_   :				List[str]								= num_samples * [prompt]
        lowerCAmelCase_   :				Optional[Any]								= pipeline.prepare_inputs(SCREAMING_SNAKE_CASE_       )

        # shard inputs and rng
        lowerCAmelCase_   :				List[Any]								= replicate(SCREAMING_SNAKE_CASE_       )
        lowerCAmelCase_   :				List[Any]								= jax.random.split(SCREAMING_SNAKE_CASE_    ,			SCREAMING_SNAKE_CASE_       )
        lowerCAmelCase_   :				List[Any]								= shard(SCREAMING_SNAKE_CASE_       )

        lowerCAmelCase_   :				Dict								= pipeline(SCREAMING_SNAKE_CASE_    ,			SCREAMING_SNAKE_CASE_    ,			SCREAMING_SNAKE_CASE_    ,			SCREAMING_SNAKE_CASE_    ,			jit=SCREAMING_SNAKE_CASE_       ).images

        assert images.shape == (num_samples, 1, 5_1_2, 5_1_2, 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(SCREAMING_SNAKE_CASE_    ,			dtype=np.floataa       ).sum() - 2_37_35_16.75)       ) < 5E-1



      def      SCREAMING_SNAKE_CASE__ (      self				:      List[Any]       ):
        lowerCAmelCase_   :				List[str]								= FlaxDDIMScheduler(
            beta_start=0.0_00_85    ,			beta_end=0.0_12    ,			beta_schedule='scaled_linear'    ,			set_alpha_to_one=SCREAMING_SNAKE_CASE_    ,			steps_offset=1    ,			)

        lowerCAmelCase_						,lowerCAmelCase_   :				Dict								= FlaxStableDiffusionPipeline.from_pretrained(
            'CompVis/stable-diffusion-v1-4'    ,			revision='bf16'    ,			dtype=jnp.bfloataa    ,			scheduler=SCREAMING_SNAKE_CASE_    ,			safety_checker=SCREAMING_SNAKE_CASE_    ,			)
        lowerCAmelCase_   :				Optional[Any]								= scheduler.create_state()

        lowerCAmelCase_   :				List[Any]								= scheduler_state

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

        lowerCAmelCase_   :				Union[str, Any]								= jax.random.PRNGKey(0       )
        lowerCAmelCase_   :				int								= 5_0

        lowerCAmelCase_   :				Tuple								= jax.device_count()
        lowerCAmelCase_   :				Any								= num_samples * [prompt]
        lowerCAmelCase_   :				Dict								= pipeline.prepare_inputs(SCREAMING_SNAKE_CASE_       )

        # shard inputs and rng
        lowerCAmelCase_   :				str								= replicate(SCREAMING_SNAKE_CASE_       )
        lowerCAmelCase_   :				Union[str, Any]								= jax.random.split(SCREAMING_SNAKE_CASE_    ,			SCREAMING_SNAKE_CASE_       )
        lowerCAmelCase_   :				str								= shard(SCREAMING_SNAKE_CASE_       )

        lowerCAmelCase_   :				Union[str, Any]								= pipeline(SCREAMING_SNAKE_CASE_    ,			SCREAMING_SNAKE_CASE_    ,			SCREAMING_SNAKE_CASE_    ,			SCREAMING_SNAKE_CASE_    ,			jit=SCREAMING_SNAKE_CASE_       ).images

        assert images.shape == (num_samples, 1, 5_1_2, 5_1_2, 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(SCREAMING_SNAKE_CASE_    ,			dtype=np.floataa       ).sum() - 2_34_76_93.5)       ) < 5E-1





      def      SCREAMING_SNAKE_CASE__ (      self				:      Optional[int]       ):
        lowerCAmelCase_   :				str								= (
            'A cinematic film still of Morgan Freeman starring as Jimi Hendrix, portrait, 40mm lens, shallow depth of'
            ' field, close up, split lighting, cinematic'
        )

        lowerCAmelCase_   :				int								= jax.device_count()
        lowerCAmelCase_   :				Optional[int]								= num_samples * [prompt]
        lowerCAmelCase_   :				int								= jax.random.split(jax.random.PRNGKey(0       )    ,			SCREAMING_SNAKE_CASE_       )

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

        lowerCAmelCase_   :				List[str]								= replicate(SCREAMING_SNAKE_CASE_       )
        lowerCAmelCase_   :				Dict								= pipeline.prepare_inputs(SCREAMING_SNAKE_CASE_       )
        lowerCAmelCase_   :				Optional[int]								= shard(SCREAMING_SNAKE_CASE_       )
        lowerCAmelCase_   :				int								= pipeline(SCREAMING_SNAKE_CASE_    ,			SCREAMING_SNAKE_CASE_    ,			SCREAMING_SNAKE_CASE_    ,			jit=SCREAMING_SNAKE_CASE_       ).images
        assert images.shape == (num_samples, 1, 5_1_2, 5_1_2, 3)
        lowerCAmelCase_   :				Optional[Any]								= images[2, 0, 2_5_6, 1_0:1_7, 1]

        # With memory efficient attention
        lowerCAmelCase_						,lowerCAmelCase_   :				Union[str, Any]								= FlaxStableDiffusionPipeline.from_pretrained(
            'CompVis/stable-diffusion-v1-4'    ,			revision='bf16'    ,			dtype=jnp.bfloataa    ,			safety_checker=SCREAMING_SNAKE_CASE_    ,			use_memory_efficient_attention=SCREAMING_SNAKE_CASE_    ,			)

        lowerCAmelCase_   :				str								= replicate(SCREAMING_SNAKE_CASE_       )
        lowerCAmelCase_   :				Tuple								= pipeline.prepare_inputs(SCREAMING_SNAKE_CASE_       )
        lowerCAmelCase_   :				Optional[int]								= shard(SCREAMING_SNAKE_CASE_       )
        lowerCAmelCase_   :				Optional[int]								= pipeline(SCREAMING_SNAKE_CASE_    ,			SCREAMING_SNAKE_CASE_    ,			SCREAMING_SNAKE_CASE_    ,			jit=SCREAMING_SNAKE_CASE_       ).images
        assert images_eff.shape == (num_samples, 1, 5_1_2, 5_1_2, 3)
        lowerCAmelCase_   :				int								= images[2, 0, 2_5_6, 1_0:1_7, 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"""







import inspect
import os
import unittest

import torch

import accelerate
from accelerate import Accelerator
from accelerate.test_utils import execute_subprocess_async, require_multi_gpu
from accelerate.utils import patch_environment







class       UpperCamelCase__   (	unittest.TestCase  ):




      """simple docstring"""




      def      SCREAMING_SNAKE_CASE__ (      self				:      Optional[int]       ):
        lowerCAmelCase_   :				Union[str, Any]								= inspect.getfile(accelerate.test_utils       )
        lowerCAmelCase_   :				Dict								= os.path.sep.join(mod_file.split(os.path.sep       )[:-1] + ['scripts', 'test_script.py']       )
        lowerCAmelCase_   :				int								= os.path.sep.join(
            mod_file.split(os.path.sep       )[:-1] + ['scripts', 'test_distributed_data_loop.py']       )
        lowerCAmelCase_   :				Union[str, Any]								= os.path.sep.join(mod_file.split(os.path.sep       )[:-1] + ['scripts', 'test_ops.py']       )



      @require_multi_gpu
      def      SCREAMING_SNAKE_CASE__ (      self				:      Any       ):
        print(F"Found {torch.cuda.device_count()} devices."       )
        lowerCAmelCase_   :				Optional[int]								= ['torchrun', F"--nproc_per_node={torch.cuda.device_count()}", self.test_file_path]
        with patch_environment(omp_num_threads=1       ):
          execute_subprocess_async(SCREAMING_SNAKE_CASE_    ,			env=os.environ.copy()       )



      @require_multi_gpu
      def      SCREAMING_SNAKE_CASE__ (      self				:      Optional[Any]       ):
        print(F"Found {torch.cuda.device_count()} devices."       )
        lowerCAmelCase_   :				int								= ['torchrun', F"--nproc_per_node={torch.cuda.device_count()}", self.operation_file_path]
        print(F"Command: {cmd}"       )
        with patch_environment(omp_num_threads=1       ):
          execute_subprocess_async(SCREAMING_SNAKE_CASE_    ,			env=os.environ.copy()       )



      @require_multi_gpu
      def      SCREAMING_SNAKE_CASE__ (      self				:      Optional[int]       ):
        lowerCAmelCase_   :				List[Any]								= ['torchrun', F"--nproc_per_node={torch.cuda.device_count()}", inspect.getfile(self.__class__       )]
        with patch_environment(omp_num_threads=1       ):
          execute_subprocess_async(SCREAMING_SNAKE_CASE_    ,			env=os.environ.copy()       )





      @require_multi_gpu
      def      SCREAMING_SNAKE_CASE__ (      self				:      Union[str, Any]       ):
        print(F"Found {torch.cuda.device_count()} devices, using 2 devices only"       )
        lowerCAmelCase_   :				Any								= ['torchrun', F"--nproc_per_node={torch.cuda.device_count()}", self.data_loop_file_path]
        with patch_environment(omp_num_threads=1    ,			cuda_visible_devices='0,1'       ):
          execute_subprocess_async(SCREAMING_SNAKE_CASE_    ,			env=os.environ.copy()       )


if __name__ == "__main__":
     lowercase__		:				Dict          =	Accelerator()
     lowercase__		:				List[Any]          =	(accelerator.state.process_index + 2, 1_0)
     lowercase__		:				Any          =	torch.randint(0, 1_0, shape).to(accelerator.device)

     lowercase__		:				List[Any]          =	""""""

     lowercase__		:				Union[str, Any]          =	accelerator.pad_across_processes(tensor)
     if tensora.shape[0] != accelerator.state.num_processes + 1:
          error_msg += f"Found shape {tensora.shape} but should have {accelerator.state.num_processes + 1} at dim 0."
     if not torch.equal(tensora[: accelerator.state.process_index + 2], tensor):
          error_msg += "Tensors have different values."
     if not torch.all(tensora[accelerator.state.process_index + 2 :] == 0):
          error_msg += "Padding was not done with the right value (0)."

     lowercase__		:				Optional[Any]          =	accelerator.pad_across_processes(tensor, pad_first=True)
     if tensora.shape[0] != accelerator.state.num_processes + 1:
          error_msg += f"Found shape {tensora.shape} but should have {accelerator.state.num_processes + 1} at dim 0."
     lowercase__		:				Any          =	accelerator.state.num_processes - accelerator.state.process_index - 1
     if not torch.equal(tensora[index:], tensor):
          error_msg += "Tensors have different values."
     if not torch.all(tensora[:index] == 0):
          error_msg += "Padding was not done with the right value (0)."

     # Raise error at the end to make sure we don't stop at the first failure.
     if len(error_msg) > 0:
          raise ValueError(error_msg)





