

import inspect
import os

import torch
from transformers import AutoModel
from transformers.testing_utils import mockenv_context
from transformers.trainer_utils import set_seed

import accelerate
from accelerate.accelerator import Accelerator
from accelerate.state import AcceleratorState
from accelerate.test_utils.testing import (
    AccelerateTestCase,
    TempDirTestCase,
    execute_subprocess_async,
    require_cuda,
    require_fsdp,
    require_multi_gpu,
    slow,
)
from accelerate.utils.constants import (
    FSDP_AUTO_WRAP_POLICY,
    FSDP_BACKWARD_PREFETCH,
    FSDP_SHARDING_STRATEGY,
    FSDP_STATE_DICT_TYPE,
)
from accelerate.utils.dataclasses import FullyShardedDataParallelPlugin
from accelerate.utils.other import patch_environment


set_seed(42)

a_     :    Union[str, Any]     		=   "bert-base-cased"
a_     :    Optional[Any]     		=   "fp16"
a_     :    Tuple     		=   "bf16"
a_     :    Dict     		=   [FPaa, BFaa]




@require_fsdp
@require_cuda
class   UpperCamelCase			(  SCREAMING_SNAKE_CASE ):

    def UpperCamelCase      (      self			:   Tuple     ):

     """simple docstring"""






     super().setUp()

     SCREAMING_SNAKE_CASE		    =      dict(
         ACCELERATE_USE_FSDP='true'				,     MASTER_ADDR='localhost'				,     MASTER_PORT='10999'				,     RANK='0'				,     LOCAL_RANK='0'				,     WORLD_SIZE='1'				,     )

    def UpperCamelCase      (      self			:   str     ):

     """simple docstring"""






     from torch.distributed.fsdp.fully_sharded_data_parallel import ShardingStrategy

     for i, strategy in enumerate(snake_case__     ):
      SCREAMING_SNAKE_CASE		    =      self.dist_env.copy()
      SCREAMING_SNAKE_CASE		    =      F"""{i + 1}"""
      SCREAMING_SNAKE_CASE		    =      strategy
      with mockenv_context(**snake_case__     ):
       SCREAMING_SNAKE_CASE		    =      FullyShardedDataParallelPlugin()
       self.assertEqual(fsdp_plugin.sharding_strategy				,     ShardingStrategy(i + 1     )     )

    def UpperCamelCase      (      self			:   int     ):

     """simple docstring"""






     from torch.distributed.fsdp.fully_sharded_data_parallel import BackwardPrefetch

     for i, prefetch_policy in enumerate(snake_case__     ):
      SCREAMING_SNAKE_CASE		    =      self.dist_env.copy()
      SCREAMING_SNAKE_CASE		    =      prefetch_policy
      with mockenv_context(**snake_case__     ):
       SCREAMING_SNAKE_CASE		    =      FullyShardedDataParallelPlugin()
       if prefetch_policy == "NO_PREFETCH":
        self.assertIsNone(fsdp_plugin.backward_prefetch     )
       else:
        self.assertEqual(fsdp_plugin.backward_prefetch				,     BackwardPrefetch(i + 1     )     )

    def UpperCamelCase      (      self			:   Dict     ):

     """simple docstring"""






     from torch.distributed.fsdp.fully_sharded_data_parallel import StateDictType

     for i, state_dict_type in enumerate(snake_case__     ):
      SCREAMING_SNAKE_CASE		    =      self.dist_env.copy()
      SCREAMING_SNAKE_CASE		    =      state_dict_type
      with mockenv_context(**snake_case__     ):
       SCREAMING_SNAKE_CASE		    =      FullyShardedDataParallelPlugin()
       self.assertEqual(fsdp_plugin.state_dict_type				,     StateDictType(i + 1     )     )
       if state_dict_type == "FULL_STATE_DICT":
        self.assertTrue(fsdp_plugin.state_dict_config.offload_to_cpu     )
        self.assertTrue(fsdp_plugin.state_dict_config.ranka_only     )

    def UpperCamelCase      (      self			:   Optional[Any]     ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE		    =      AutoModel.from_pretrained(snake_case__     )
     for policy in FSDP_AUTO_WRAP_POLICY:
      SCREAMING_SNAKE_CASE		    =      self.dist_env.copy()
      SCREAMING_SNAKE_CASE		    =      policy
      if policy == "TRANSFORMER_BASED_WRAP":
       SCREAMING_SNAKE_CASE		    =      'BertLayer'
      elif policy == "SIZE_BASED_WRAP":
       SCREAMING_SNAKE_CASE		    =      '2000'
      with mockenv_context(**snake_case__     ):
       SCREAMING_SNAKE_CASE		    =      FullyShardedDataParallelPlugin()
       fsdp_plugin.set_auto_wrap_policy(snake_case__     )
       if policy == "NO_WRAP":
        self.assertIsNone(fsdp_plugin.auto_wrap_policy     )
       else:
        self.assertIsNotNone(fsdp_plugin.auto_wrap_policy     )

     SCREAMING_SNAKE_CASE		    =      self.dist_env.copy()
     SCREAMING_SNAKE_CASE		    =      'TRANSFORMER_BASED_WRAP'
     SCREAMING_SNAKE_CASE		    =      'T5Layer'
     with mockenv_context(**snake_case__     ):
      SCREAMING_SNAKE_CASE		    =      FullyShardedDataParallelPlugin()
      with self.assertRaises(snake_case__     ) as cm:
       fsdp_plugin.set_auto_wrap_policy(snake_case__     )
      self.assertTrue('Could not find the transformer layer class to wrap in the model.' in str(cm.exception     )     )

     SCREAMING_SNAKE_CASE		    =      self.dist_env.copy()
     SCREAMING_SNAKE_CASE		    =      'SIZE_BASED_WRAP'
     SCREAMING_SNAKE_CASE		    =      '0'
     with mockenv_context(**snake_case__     ):
      SCREAMING_SNAKE_CASE		    =      FullyShardedDataParallelPlugin()
      fsdp_plugin.set_auto_wrap_policy(snake_case__     )
      self.assertIsNone(fsdp_plugin.auto_wrap_policy     )

    def UpperCamelCase      (      self			:   int     ):

     """simple docstring"""






     from torch.distributed.fsdp.fully_sharded_data_parallel import MixedPrecision
     from torch.distributed.fsdp.sharded_grad_scaler import ShardedGradScaler

     for mp_dtype in dtypes:
      SCREAMING_SNAKE_CASE		    =      self.dist_env.copy()
      SCREAMING_SNAKE_CASE		    =      mp_dtype
      with mockenv_context(**snake_case__     ):
       SCREAMING_SNAKE_CASE		    =      Accelerator()
       if mp_dtype == "fp16":
        SCREAMING_SNAKE_CASE		    =      torch.floataa
       elif mp_dtype == "bf16":
        SCREAMING_SNAKE_CASE		    =      torch.bfloataa
       SCREAMING_SNAKE_CASE		    =      MixedPrecision(param_dtype=snake_case__				,     reduce_dtype=snake_case__				,     buffer_dtype=snake_case__     )
       self.assertEqual(accelerator.state.fsdp_plugin.mixed_precision_policy				,     snake_case__     )
       if mp_dtype == FPaa:
        self.assertTrue(isinstance(accelerator.scaler				,     snake_case__     )     )
       elif mp_dtype == BFaa:
        self.assertIsNone(accelerator.scaler     )
       AcceleratorState._reset_state(snake_case__     )




    def UpperCamelCase      (      self			:   Optional[int]     ):

     """simple docstring"""






     from torch.distributed.fsdp.fully_sharded_data_parallel import CPUOffload

     for flag in [True, False]:
      SCREAMING_SNAKE_CASE		    =      self.dist_env.copy()
      SCREAMING_SNAKE_CASE		    =      str(snake_case__     ).lower()
      with mockenv_context(**snake_case__     ):
       SCREAMING_SNAKE_CASE		    =      FullyShardedDataParallelPlugin()
       self.assertEqual(fsdp_plugin.cpu_offload				,     CPUOffload(offload_params=snake_case__     )     )




@require_fsdp
@require_multi_gpu
@slow
class   UpperCamelCase			(  SCREAMING_SNAKE_CASE ):

    def UpperCamelCase      (      self			:   Any     ):

     """simple docstring"""






     super().setUp()
     SCREAMING_SNAKE_CASE		    =      0.82
     SCREAMING_SNAKE_CASE		    =      [
         'fsdp_shard_grad_op_transformer_based_wrap',
         'fsdp_full_shard_transformer_based_wrap',
     ]
     SCREAMING_SNAKE_CASE		    =      {
         'multi_gpu_fp16': 3_2_0_0,
         'fsdp_shard_grad_op_transformer_based_wrap_fp16': 2_0_0_0,
         'fsdp_full_shard_transformer_based_wrap_fp16': 1_9_0_0,
         # Disabling below test as it overwhelms the RAM memory usage
         # on CI self-hosted runner leading to tests getting killed.
         # "fsdp_full_shard_cpu_offload_transformer_based_wrap_fp32": 1500,  # fp16 was leading to indefinite hang
     }
     SCREAMING_SNAKE_CASE		    =      1_6_0
     SCREAMING_SNAKE_CASE		    =      1_6_0

     SCREAMING_SNAKE_CASE		    =      inspect.getfile(accelerate.test_utils     )
     SCREAMING_SNAKE_CASE		    =      os.path.sep.join(mod_file.split(os.path.sep     )[:-1] + ['scripts', 'external_deps']     )

    def UpperCamelCase      (      self			:   Any     ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE		    =      os.path.join(self.test_scripts_folder				,     'test_performance.py'     )
     SCREAMING_SNAKE_CASE		    =      ['accelerate', 'launch', '--num_processes=2', '--num_machines=1', '--machine_rank=0', '--use_fsdp']
     for config in self.performance_configs:
      SCREAMING_SNAKE_CASE		    =      cmd.copy()
      for i, strategy in enumerate(snake_case__     ):
       if strategy.lower() in config:
        cmd_config.append(F"""--fsdp_sharding_strategy={i+1}"""     )
        break

      if "fp32" in config:
       cmd_config.append('--mixed_precision=no'     )
      else:
       cmd_config.append('--mixed_precision=fp16'     )

      if "cpu_offload" in config:
       cmd_config.append('--fsdp_offload_params=True'     )

      for policy in FSDP_AUTO_WRAP_POLICY:
       if policy.lower() in config:
        cmd_config.append(F"""--fsdp_auto_wrap_policy={policy}"""     )
        break

      if policy == "TRANSFORMER_BASED_WRAP":
       cmd_config.append('--fsdp_transformer_layer_cls_to_wrap=BertLayer'     )
      elif policy == "SIZE_BASED_WRAP":
       cmd_config.append('--fsdp_min_num_params=2000'     )

      cmd_config.extend(
          [
              self.test_file_path,
              F"""--output_dir={self.tmpdir}""",
              F"""--performance_lower_bound={self.performance_lower_bound}""",
          ]     )
      with patch_environment(omp_num_threads=1     ):
       execute_subprocess_async(snake_case__				,     env=os.environ.copy()     )

    def UpperCamelCase      (      self			:   str     ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE		    =      os.path.join(self.test_scripts_folder				,     'test_checkpointing.py'     )
     SCREAMING_SNAKE_CASE		    =      [
         'accelerate',
         'launch',
         '--num_processes=2',
         '--num_machines=1',
         '--machine_rank=0',
         '--use_fsdp',
         '--mixed_precision=fp16',
         '--fsdp_transformer_layer_cls_to_wrap=BertLayer',
     ]

     for i, strategy in enumerate(snake_case__     ):
      SCREAMING_SNAKE_CASE		    =      cmd.copy()
      cmd_config.append(F"""--fsdp_sharding_strategy={i+1}"""     )
      if strategy != "FULL_SHARD":
       continue
      SCREAMING_SNAKE_CASE		    =      len(snake_case__     )
      for state_dict_type in FSDP_STATE_DICT_TYPE:
       SCREAMING_SNAKE_CASE		    =      cmd_config[:state_dict_config_index]
       cmd_config.append(F"""--fsdp_state_dict_type={state_dict_type}"""     )
       cmd_config.extend(
           [
               self.test_file_path,
               F"""--output_dir={self.tmpdir}""",
               '--partial_train_epoch=1',
           ]     )
       with patch_environment(omp_num_threads=1     ):
        execute_subprocess_async(snake_case__				,     env=os.environ.copy()     )

       SCREAMING_SNAKE_CASE		    =      cmd_config[:-1]
       SCREAMING_SNAKE_CASE		    =      os.path.join(self.tmpdir				,     'epoch_0'     )
       cmd_config.extend(
           [
               F"""--resume_from_checkpoint={resume_from_checkpoint}""",
           ]     )
       with patch_environment(omp_num_threads=1     ):
        execute_subprocess_async(snake_case__				,     env=os.environ.copy()     )




    def UpperCamelCase      (      self			:   Dict     ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE		    =      os.path.join(self.test_scripts_folder				,     'test_peak_memory_usage.py'     )
     SCREAMING_SNAKE_CASE		    =      [
         'accelerate',
         'launch',
         '--num_processes=2',
         '--num_machines=1',
         '--machine_rank=0',
     ]
     for spec, peak_mem_upper_bound in self.peak_memory_usage_upper_bound.items():
      SCREAMING_SNAKE_CASE		    =      cmd.copy()
      if "fp16" in spec:
       cmd_config.extend(['--mixed_precision=fp16']     )
      else:
       cmd_config.extend(['--mixed_precision=no']     )

      if "multi_gpu" in spec:
       continue
      else:
       cmd_config.extend(['--use_fsdp']     )
       for i, strategy in enumerate(snake_case__     ):
        if strategy.lower() in spec:
         cmd_config.append(F"""--fsdp_sharding_strategy={i+1}"""     )
         break

       if "cpu_offload" in spec:
        cmd_config.append('--fsdp_offload_params=True'     )

       for policy in FSDP_AUTO_WRAP_POLICY:
        if policy.lower() in spec:
         cmd_config.append(F"""--fsdp_auto_wrap_policy={policy}"""     )
         break

       if policy == "TRANSFORMER_BASED_WRAP":
        cmd_config.append('--fsdp_transformer_layer_cls_to_wrap=BertLayer'     )
       elif policy == "SIZE_BASED_WRAP":
        cmd_config.append('--fsdp_min_num_params=2000'     )

      cmd_config.extend(
          [
              self.test_file_path,
              F"""--output_dir={self.tmpdir}""",
              F"""--peak_memory_upper_bound={peak_mem_upper_bound}""",
              F"""--n_train={self.n_train}""",
              F"""--n_val={self.n_val}""",
          ]     )
      with patch_environment(omp_num_threads=1     ):
       execute_subprocess_async(snake_case__				,     env=os.environ.copy()     )




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

    def UpperCamelCase      (      self			:   Tuple     ):

     """simple docstring"""






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

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

    def UpperCamelCase      (      self			:   Optional[Any]     ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE    ,						SCREAMING_SNAKE_CASE		    =      FlaxStableDiffusionPipeline.from_pretrained(
         'hf-internal-testing/tiny-stable-diffusion-pipe'				,     safety_checker=snake_case__     )

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

     SCREAMING_SNAKE_CASE		    =      jax.random.PRNGKey(0     )
     SCREAMING_SNAKE_CASE		    =      4

     SCREAMING_SNAKE_CASE		    =      jax.device_count()
     SCREAMING_SNAKE_CASE		    =      num_samples * [prompt]
     SCREAMING_SNAKE_CASE		    =      pipeline.prepare_inputs(snake_case__     )

     # shard inputs and rng
     SCREAMING_SNAKE_CASE		    =      replicate(snake_case__     )
     SCREAMING_SNAKE_CASE		    =      jax.random.split(snake_case__				,     snake_case__     )
     SCREAMING_SNAKE_CASE		    =      shard(snake_case__     )

     SCREAMING_SNAKE_CASE		    =      pipeline(snake_case__				,     snake_case__				,     snake_case__				,     snake_case__				,     jit=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_514_745     ) < 1E-3
      assert np.abs(np.abs(snake_case__				,     dtype=np.floataa     ).sum() - 49_947.875     ) < 5E-1

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

    def UpperCamelCase      (      self			:   Union[str, Any]     ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE    ,						SCREAMING_SNAKE_CASE		    =      FlaxStableDiffusionPipeline.from_pretrained(
         'CompVis/stable-diffusion-v1-4'				,     revision='flax'				,     safety_checker=snake_case__     )

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

     SCREAMING_SNAKE_CASE		    =      jax.random.PRNGKey(0     )
     SCREAMING_SNAKE_CASE		    =      5_0

     SCREAMING_SNAKE_CASE		    =      jax.device_count()
     SCREAMING_SNAKE_CASE		    =      num_samples * [prompt]
     SCREAMING_SNAKE_CASE		    =      pipeline.prepare_inputs(snake_case__     )

     # shard inputs and rng
     SCREAMING_SNAKE_CASE		    =      replicate(snake_case__     )
     SCREAMING_SNAKE_CASE		    =      jax.random.split(snake_case__				,     snake_case__     )
     SCREAMING_SNAKE_CASE		    =      shard(snake_case__     )

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

    def UpperCamelCase      (      self			:   Optional[Any]     ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE    ,						SCREAMING_SNAKE_CASE		    =      FlaxStableDiffusionPipeline.from_pretrained(
         'CompVis/stable-diffusion-v1-4'				,     revision='bf16'				,     dtype=jnp.bfloataa				,     safety_checker=snake_case__     )

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

     SCREAMING_SNAKE_CASE		    =      jax.random.PRNGKey(0     )
     SCREAMING_SNAKE_CASE		    =      5_0

     SCREAMING_SNAKE_CASE		    =      jax.device_count()
     SCREAMING_SNAKE_CASE		    =      num_samples * [prompt]
     SCREAMING_SNAKE_CASE		    =      pipeline.prepare_inputs(snake_case__     )

     # shard inputs and rng
     SCREAMING_SNAKE_CASE		    =      replicate(snake_case__     )
     SCREAMING_SNAKE_CASE		    =      jax.random.split(snake_case__				,     snake_case__     )
     SCREAMING_SNAKE_CASE		    =      shard(snake_case__     )

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

    def UpperCamelCase      (      self			:   List[Any]     ):

     """simple docstring"""






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

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

     SCREAMING_SNAKE_CASE		    =      jax.random.PRNGKey(0     )
     SCREAMING_SNAKE_CASE		    =      5_0

     SCREAMING_SNAKE_CASE		    =      jax.device_count()
     SCREAMING_SNAKE_CASE		    =      num_samples * [prompt]
     SCREAMING_SNAKE_CASE		    =      pipeline.prepare_inputs(snake_case__     )

     # shard inputs and rng
     SCREAMING_SNAKE_CASE		    =      replicate(snake_case__     )
     SCREAMING_SNAKE_CASE		    =      jax.random.split(snake_case__				,     snake_case__     )
     SCREAMING_SNAKE_CASE		    =      shard(snake_case__     )

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

    def UpperCamelCase      (      self			:   Dict     ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE		    =      FlaxDDIMScheduler(
         beta_start=0.00_085				,     beta_end=0.012				,     beta_schedule='scaled_linear'				,     set_alpha_to_one=snake_case__				,     steps_offset=1				,     )

     SCREAMING_SNAKE_CASE    ,						SCREAMING_SNAKE_CASE		    =      FlaxStableDiffusionPipeline.from_pretrained(
         'CompVis/stable-diffusion-v1-4'				,     revision='bf16'				,     dtype=jnp.bfloataa				,     scheduler=snake_case__				,     safety_checker=snake_case__				,     )
     SCREAMING_SNAKE_CASE		    =      scheduler.create_state()

     SCREAMING_SNAKE_CASE		    =      scheduler_state

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

     SCREAMING_SNAKE_CASE		    =      jax.random.PRNGKey(0     )
     SCREAMING_SNAKE_CASE		    =      5_0

     SCREAMING_SNAKE_CASE		    =      jax.device_count()
     SCREAMING_SNAKE_CASE		    =      num_samples * [prompt]
     SCREAMING_SNAKE_CASE		    =      pipeline.prepare_inputs(snake_case__     )

     # shard inputs and rng
     SCREAMING_SNAKE_CASE		    =      replicate(snake_case__     )
     SCREAMING_SNAKE_CASE		    =      jax.random.split(snake_case__				,     snake_case__     )
     SCREAMING_SNAKE_CASE		    =      shard(snake_case__     )

     SCREAMING_SNAKE_CASE		    =      pipeline(snake_case__				,     snake_case__				,     snake_case__				,     snake_case__				,     jit=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.045_043_945)     ) < 1E-3
      assert np.abs((np.abs(snake_case__				,     dtype=np.floataa     ).sum() - 2_347_693.5)     ) < 5E-1




    def UpperCamelCase      (      self			:   Tuple     ):

     """simple docstring"""






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

     SCREAMING_SNAKE_CASE		    =      jax.device_count()
     SCREAMING_SNAKE_CASE		    =      num_samples * [prompt]
     SCREAMING_SNAKE_CASE		    =      jax.random.split(jax.random.PRNGKey(0     )				,     snake_case__     )

     SCREAMING_SNAKE_CASE    ,						SCREAMING_SNAKE_CASE		    =      FlaxStableDiffusionPipeline.from_pretrained(
         'CompVis/stable-diffusion-v1-4'				,     revision='bf16'				,     dtype=jnp.bfloataa				,     safety_checker=snake_case__				,     )

     SCREAMING_SNAKE_CASE		    =      replicate(snake_case__     )
     SCREAMING_SNAKE_CASE		    =      pipeline.prepare_inputs(snake_case__     )
     SCREAMING_SNAKE_CASE		    =      shard(snake_case__     )
     SCREAMING_SNAKE_CASE		    =      pipeline(snake_case__				,     snake_case__				,     snake_case__				,     jit=snake_case__     ).images
     assert images.shape == (num_samples, 1, 5_1_2, 5_1_2, 3)
     SCREAMING_SNAKE_CASE		    =      images[2, 0, 2_5_6, 1_0:1_7, 1]

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

     SCREAMING_SNAKE_CASE		    =      replicate(snake_case__     )
     SCREAMING_SNAKE_CASE		    =      pipeline.prepare_inputs(snake_case__     )
     SCREAMING_SNAKE_CASE		    =      shard(snake_case__     )
     SCREAMING_SNAKE_CASE		    =      pipeline(snake_case__				,     snake_case__				,     snake_case__				,     jit=snake_case__     ).images
     assert images_eff.shape == (num_samples, 1, 5_1_2, 5_1_2, 3)
     SCREAMING_SNAKE_CASE		    =      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

