


from typing import Optional, Tuple, Union

import flax
import flax.linen as nn
import jax
import jax.numpy as jnp
from flax.core.frozen_dict import FrozenDict

from ..configuration_utils import ConfigMixin, flax_register_to_config
from ..utils import BaseOutput
from .embeddings_flax import FlaxTimestepEmbedding, FlaxTimesteps
from .modeling_flax_utils import FlaxModelMixin
from .unet_ad_blocks_flax import (
    FlaxCrossAttnDownBlockaD,
    FlaxDownBlockaD,
    FlaxUNetMidBlockaDCrossAttn,
)





@flax.struct.dataclass
class        lowerCAmelCase__ (					__lowerCamelCase  ):






						"""simple docstring"""
						__UpperCAmelCase       :						List[Any]    	=			42
						__UpperCAmelCase       :						Optional[int]    	=			42





class        lowerCAmelCase__ (					nn.Module  ):






						"""simple docstring"""
						__UpperCAmelCase       :						List[Any]    	=			42
						__UpperCAmelCase       :						Dict    	=			(16, 32, 96, 256)
						__UpperCAmelCase       :						Optional[int]    	=			jnp.floataa




						def 	_UpperCamelCase					(       self      ):
							lowerCamelCase_					:				str              =  nn.Conv(
							    self.block_out_channels[0] ,				kernel_size=(3, 3) ,				padding=((1, 1), (1, 1)) ,				dtype=self.dtype ,				)

							lowerCamelCase_					:				Any              =  []
							for i in range(len(self.block_out_channels      ) - 1      ):
								lowerCamelCase_					:				Optional[int]              =  self.block_out_channels[i]
								lowerCamelCase_					:				int              =  self.block_out_channels[i + 1]
								lowerCamelCase_					:				str              =  nn.Conv(
								    snake_case__ ,				kernel_size=(3, 3) ,				padding=((1, 1), (1, 1)) ,				dtype=self.dtype ,				)
								blocks.append(snake_case__      )
								lowerCamelCase_					:				Optional[int]              =  nn.Conv(
								    snake_case__ ,				kernel_size=(3, 3) ,				strides=(2, 2) ,				padding=((1, 1), (1, 1)) ,				dtype=self.dtype ,				)
								blocks.append(snake_case__      )
							lowerCamelCase_					:				Optional[Any]              =  blocks

							lowerCamelCase_					:				Optional[Any]              =  nn.Conv(
							    self.conditioning_embedding_channels ,				kernel_size=(3, 3) ,				padding=((1, 1), (1, 1)) ,				kernel_init=nn.initializers.zeros_init() ,				bias_init=nn.initializers.zeros_init() ,				dtype=self.dtype ,				)




						def __call__(       self ,				a_      ):
							lowerCamelCase_					:				List[Any]              =  self.conv_in(snake_case__      )
							lowerCamelCase_					:				Dict              =  nn.silu(snake_case__      )

							for block in self.blocks:
								lowerCamelCase_					:				Union[str, Any]              =  block(snake_case__      )
								lowerCamelCase_					:				int              =  nn.silu(snake_case__      )

							lowerCamelCase_					:				Any              =  self.conv_out(snake_case__      )

							return embedding





@flax_register_to_config
class        lowerCAmelCase__ (					nn.Module,							__lowerCamelCase,							__lowerCamelCase  ):






						"""simple docstring"""
						__UpperCAmelCase       :						Optional[Any]    	=			32
						__UpperCAmelCase       :						Union[str, Any]    	=			4
						__UpperCAmelCase       :						Optional[Any]    	=			(
						    '''CrossAttnDownBlock2D''',
						    '''CrossAttnDownBlock2D''',
						    '''CrossAttnDownBlock2D''',
						    '''DownBlock2D''',
						)
						__UpperCAmelCase       :						Union[str, Any]    	=			False
						__UpperCAmelCase       :						List[Any]    	=			(320, 640, 1280, 1280)
						__UpperCAmelCase       :						Optional[Any]    	=			2
						__UpperCAmelCase       :						Tuple    	=			8
						__UpperCAmelCase       :						Union[str, Any]    	=			None
						__UpperCAmelCase       :						List[str]    	=			1280
						__UpperCAmelCase       :						List[str]    	=			0.0
						__UpperCAmelCase       :						List[Any]    	=			False
						__UpperCAmelCase       :						Optional[Any]    	=			jnp.floataa
						__UpperCAmelCase       :						Any    	=			True
						__UpperCAmelCase       :						Dict    	=			0
						__UpperCAmelCase       :						List[str]    	=			'''rgb'''
						__UpperCAmelCase       :						Dict    	=			(16, 32, 96, 256)




						def 	_UpperCamelCase					(       self ,				a_      ):
							lowerCamelCase_					:				Optional[Any]              =  (1, self.in_channels, self.sample_size, self.sample_size)
							lowerCamelCase_					:				Optional[int]              =  jnp.zeros(snake_case__ ,				dtype=jnp.floataa      )
							lowerCamelCase_					:				Dict              =  jnp.ones((1,) ,				dtype=jnp.intaa      )
							lowerCamelCase_					:				List[str]              =  jnp.zeros((1, 1, self.cross_attention_dim) ,				dtype=jnp.floataa      )
							lowerCamelCase_					:				Tuple              =  (1, 3, self.sample_size * 8, self.sample_size * 8)
							lowerCamelCase_					:				str              =  jnp.zeros(snake_case__ ,				dtype=jnp.floataa      )

							lowerCamelCase_       ,lowerCamelCase_					:				Any              =  jax.random.split(snake_case__      )
							lowerCamelCase_					:				Tuple              =  {"params": params_rng, "dropout": dropout_rng}

							return self.init(snake_case__ ,				snake_case__ ,				snake_case__ ,				snake_case__ ,				snake_case__      )["params"]




						def 	_UpperCamelCase					(       self      ):
							lowerCamelCase_					:				List[str]              =  self.block_out_channels
							lowerCamelCase_					:				List[str]              =  block_out_channels[0] * 4

							# If `num_attention_heads` is not defined (which is the case for most models)
							# it will default to `attention_head_dim`. This looks weird upon first reading it and it is.
							# The reason for this behavior is to correct for incorrectly named variables that were introduced
							# when this library was created. The incorrect naming was only discovered much later in https://github.com/huggingface/diffusers/issues/2011#issuecomment-1547958131
							# Changing `attention_head_dim` to `num_attention_heads` for 40,000+ configurations is too backwards breaking
							# which is why we correct for the naming here.
							lowerCamelCase_					:				Optional[int]              =  self.num_attention_heads or self.attention_head_dim

							# input
							lowerCamelCase_					:				Union[str, Any]              =  nn.Conv(
							    block_out_channels[0] ,				kernel_size=(3, 3) ,				strides=(1, 1) ,				padding=((1, 1), (1, 1)) ,				dtype=self.dtype ,				)

							# time
							lowerCamelCase_					:				Any              =  FlaxTimesteps(
							    block_out_channels[0] ,				flip_sin_to_cos=self.flip_sin_to_cos ,				freq_shift=self.config.freq_shift      )
							lowerCamelCase_					:				Any              =  FlaxTimestepEmbedding(snake_case__ ,				dtype=self.dtype      )

							lowerCamelCase_					:				Any              =  FlaxControlNetConditioningEmbedding(
							    conditioning_embedding_channels=block_out_channels[0] ,				block_out_channels=self.conditioning_embedding_out_channels ,				)

							lowerCamelCase_					:				Any              =  self.only_cross_attention
							if isinstance(snake_case__ ,				snake_case__      ):
								lowerCamelCase_					:				int              =  (only_cross_attention,) * len(self.down_block_types      )

							if isinstance(snake_case__ ,				snake_case__      ):
								lowerCamelCase_					:				Optional[int]              =  (num_attention_heads,) * len(self.down_block_types      )

							# down
							lowerCamelCase_					:				int              =  []
							lowerCamelCase_					:				Dict              =  []

							lowerCamelCase_					:				str              =  block_out_channels[0]

							lowerCamelCase_					:				str              =  nn.Conv(
							    snake_case__ ,				kernel_size=(1, 1) ,				padding="VALID" ,				kernel_init=nn.initializers.zeros_init() ,				bias_init=nn.initializers.zeros_init() ,				dtype=self.dtype ,				)
							controlnet_down_blocks.append(snake_case__      )

							for i, down_block_type in enumerate(self.down_block_types      ):
								lowerCamelCase_					:				Any              =  output_channel
								lowerCamelCase_					:				int              =  block_out_channels[i]
								lowerCamelCase_					:				Tuple              =  i == len(snake_case__      ) - 1

								if down_block_type == "CrossAttnDownBlock2D":
									lowerCamelCase_					:				Any              =  FlaxCrossAttnDownBlockaD(
									    in_channels=snake_case__ ,				out_channels=snake_case__ ,				dropout=self.dropout ,				num_layers=self.layers_per_block ,				num_attention_heads=num_attention_heads[i] ,				add_downsample=not is_final_block ,				use_linear_projection=self.use_linear_projection ,				only_cross_attention=only_cross_attention[i] ,				dtype=self.dtype ,				)
								else:
									lowerCamelCase_					:				Optional[int]              =  FlaxDownBlockaD(
									    in_channels=snake_case__ ,				out_channels=snake_case__ ,				dropout=self.dropout ,				num_layers=self.layers_per_block ,				add_downsample=not is_final_block ,				dtype=self.dtype ,				)

								down_blocks.append(snake_case__      )

								for _ in range(self.layers_per_block      ):
									lowerCamelCase_					:				List[Any]              =  nn.Conv(
									    snake_case__ ,				kernel_size=(1, 1) ,				padding="VALID" ,				kernel_init=nn.initializers.zeros_init() ,				bias_init=nn.initializers.zeros_init() ,				dtype=self.dtype ,				)
									controlnet_down_blocks.append(snake_case__      )

								if not is_final_block:
									lowerCamelCase_					:				Optional[int]              =  nn.Conv(
									    snake_case__ ,				kernel_size=(1, 1) ,				padding="VALID" ,				kernel_init=nn.initializers.zeros_init() ,				bias_init=nn.initializers.zeros_init() ,				dtype=self.dtype ,				)
									controlnet_down_blocks.append(snake_case__      )

							lowerCamelCase_					:				str              =  down_blocks
							lowerCamelCase_					:				str              =  controlnet_down_blocks

							# mid
							lowerCamelCase_					:				Any              =  block_out_channels[-1]
							lowerCamelCase_					:				Union[str, Any]              =  FlaxUNetMidBlockaDCrossAttn(
							    in_channels=snake_case__ ,				dropout=self.dropout ,				num_attention_heads=num_attention_heads[-1] ,				use_linear_projection=self.use_linear_projection ,				dtype=self.dtype ,				)

							lowerCamelCase_					:				Dict              =  nn.Conv(
							    snake_case__ ,				kernel_size=(1, 1) ,				padding="VALID" ,				kernel_init=nn.initializers.zeros_init() ,				bias_init=nn.initializers.zeros_init() ,				dtype=self.dtype ,				)




						def __call__(       self ,				a_ ,				a_ ,				a_ ,				a_ ,				a_ = 1.0 ,				a_ = True ,				a_ = False ,				):
							lowerCamelCase_					:				int              =  self.controlnet_conditioning_channel_order
							if channel_order == "bgr":
								lowerCamelCase_					:				str              =  jnp.flip(snake_case__ ,				axis=1      )

							# 1. time
							if not isinstance(snake_case__ ,				jnp.ndarray      ):
								lowerCamelCase_					:				Any              =  jnp.array([timesteps] ,				dtype=jnp.intaa      )
							elif isinstance(snake_case__ ,				jnp.ndarray      ) and len(timesteps.shape      ) == 0:
								lowerCamelCase_					:				Union[str, Any]              =  timesteps.astype(dtype=jnp.floataa      )
								lowerCamelCase_					:				int              =  jnp.expand_dims(snake_case__ ,				0      )

							lowerCamelCase_					:				Optional[Any]              =  self.time_proj(snake_case__      )
							lowerCamelCase_					:				Optional[Any]              =  self.time_embedding(snake_case__      )

							# 2. pre-process
							lowerCamelCase_					:				List[str]              =  jnp.transpose(snake_case__ ,				(0, 2, 3, 1)      )
							lowerCamelCase_					:				Union[str, Any]              =  self.conv_in(snake_case__      )

							lowerCamelCase_					:				str              =  jnp.transpose(snake_case__ ,				(0, 2, 3, 1)      )
							lowerCamelCase_					:				Any              =  self.controlnet_cond_embedding(snake_case__      )
							sample += controlnet_cond

							# 3. down
							lowerCamelCase_					:				Tuple              =  (sample,)
							for down_block in self.down_blocks:
								if isinstance(snake_case__ ,				snake_case__      ):
									lowerCamelCase_       ,lowerCamelCase_					:				int              =  down_block(snake_case__ ,				snake_case__ ,				snake_case__ ,				deterministic=not train      )
								else:
									lowerCamelCase_       ,lowerCamelCase_					:				Any              =  down_block(snake_case__ ,				snake_case__ ,				deterministic=not train      )
								down_block_res_samples += res_samples

							# 4. mid
							lowerCamelCase_					:				List[str]              =  self.mid_block(snake_case__ ,				snake_case__ ,				snake_case__ ,				deterministic=not train      )

							# 5. contronet blocks
							lowerCamelCase_					:				Any              =  ()
							for down_block_res_sample, controlnet_block in zip(snake_case__ ,				self.controlnet_down_blocks      ):
								lowerCamelCase_					:				Any              =  controlnet_block(snake_case__      )
								controlnet_down_block_res_samples += (down_block_res_sample,)

							lowerCamelCase_					:				Dict              =  controlnet_down_block_res_samples

							lowerCamelCase_					:				Dict              =  self.controlnet_mid_block(snake_case__      )

							# 6. scaling
							lowerCamelCase_					:				Tuple              =  [sample * conditioning_scale for sample in down_block_res_samples]
							mid_block_res_sample *= conditioning_scale

							if not return_dict:
								return (down_block_res_samples, mid_block_res_sample)

							return FlaxControlNetOutput(
							    down_block_res_samples=snake_case__ ,				mid_block_res_sample=snake_case__      )




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

