
"""simple docstring"""






import argparse
import json
import os

import numpy as np
import PIL
import requests
import tensorflow.keras.applications.efficientnet as efficientnet
import torch
from huggingface_hub import hf_hub_download
from PIL import Image
from tensorflow.keras.preprocessing import image

from transformers import (
    EfficientNetConfig,
    EfficientNetForImageClassification,
    EfficientNetImageProcessor,
)
from transformers.utils import logging


logging.set_verbosity_info()
__UpperCAmelCase       =  logging.get_logger(__name__)

__UpperCAmelCase       =  {
    'b0': efficientnet.EfficientNetBa,
    'b1': efficientnet.EfficientNetBa,
    'b2': efficientnet.EfficientNetBa,
    'b3': efficientnet.EfficientNetBa,
    'b4': efficientnet.EfficientNetBa,
    'b5': efficientnet.EfficientNetBa,
    'b6': efficientnet.EfficientNetBa,
    'b7': efficientnet.EfficientNetBa,
}

__UpperCAmelCase       =  {
    'b0': {
        'hidden_dim': 12_80,
        'width_coef': 1.0,
        'depth_coef': 1.0,
        'image_size': 2_24,
        'dropout_rate': 0.2,
        'dw_padding': [],
    },
    'b1': {
        'hidden_dim': 12_80,
        'width_coef': 1.0,
        'depth_coef': 1.1,
        'image_size': 2_40,
        'dropout_rate': 0.2,
        'dw_padding': [16],
    },
    'b2': {
        'hidden_dim': 14_08,
        'width_coef': 1.1,
        'depth_coef': 1.2,
        'image_size': 2_60,
        'dropout_rate': 0.3,
        'dw_padding': [5, 8, 16],
    },
    'b3': {
        'hidden_dim': 15_36,
        'width_coef': 1.2,
        'depth_coef': 1.4,
        'image_size': 3_00,
        'dropout_rate': 0.3,
        'dw_padding': [5, 18],
    },
    'b4': {
        'hidden_dim': 17_92,
        'width_coef': 1.4,
        'depth_coef': 1.8,
        'image_size': 3_80,
        'dropout_rate': 0.4,
        'dw_padding': [6],
    },
    'b5': {
        'hidden_dim': 20_48,
        'width_coef': 1.6,
        'depth_coef': 2.2,
        'image_size': 4_56,
        'dropout_rate': 0.4,
        'dw_padding': [13, 27],
    },
    'b6': {
        'hidden_dim': 23_04,
        'width_coef': 1.8,
        'depth_coef': 2.6,
        'image_size': 5_28,
        'dropout_rate': 0.5,
        'dw_padding': [31],
    },
    'b7': {
        'hidden_dim': 25_60,
        'width_coef': 2.0,
        'depth_coef': 3.1,
        'image_size': 6_00,
        'dropout_rate': 0.5,
        'dw_padding': [18],
    },
}






def 						_snake_case       (	lowercase__		:  Tuple		)  ->				str:


			'''simple docstring'''





			lowerCAmelCase_    :Tuple			     =  EfficientNetConfig()
			lowerCAmelCase_    :str			     =  CONFIG_MAP[model_name]["""hidden_dim"""]
			lowerCAmelCase_    :List[Any]			     =  CONFIG_MAP[model_name]["""width_coef"""]
			lowerCAmelCase_    :Any			     =  CONFIG_MAP[model_name]["""depth_coef"""]
			lowerCAmelCase_    :str			     =  CONFIG_MAP[model_name]["""image_size"""]
			lowerCAmelCase_    :Optional[int]			     =  CONFIG_MAP[model_name]["""dropout_rate"""]
			lowerCAmelCase_    :Dict			     =  CONFIG_MAP[model_name]["""dw_padding"""]

			lowerCAmelCase_    :Dict			     =  """huggingface/label-files"""
			lowerCAmelCase_    :Optional[int]			     =  """imagenet-1k-id2label.json"""
			lowerCAmelCase_    :List[str]			     =  1_0_0_0
			lowerCAmelCase_    :Any			     =  json.load(open(hf_hub_download(lowercase__      ,     lowercase__      ,     repo_type="""dataset"""		)      ,     """r"""		)		)
			lowerCAmelCase_    :int			     =  {int(lowercase__		): v for k, v in idalabel.items()}

			lowerCAmelCase_    :Any			     =  idalabel
			lowerCAmelCase_    :Union[str, Any]			     =  {v: k for k, v in idalabel.items()}
			return config






def 						_snake_case       (	)  ->				int:


			'''simple docstring'''





			lowerCAmelCase_    :Dict			     =  """http://images.cocodataset.org/val2017/000000039769.jpg"""
			lowerCAmelCase_    :List[Any]			     =  Image.open(requests.get(lowercase__      ,     stream=lowercase__		).raw		)
			return im






def 						_snake_case       (	lowercase__		:  int		)  ->				Tuple:


			'''simple docstring'''





			lowerCAmelCase_    :str			     =  CONFIG_MAP[model_name]["""image_size"""]
			lowerCAmelCase_    :str			     =  EfficientNetImageProcessor(
			    size={"""height""": size, """width""": size}      ,     image_mean=[0.485, 0.456, 0.406]      ,     image_std=[0.47853944, 0.4732864, 0.47434163]      ,     do_center_crop=lowercase__      ,     )
			return preprocessor






def 						_snake_case       (	lowercase__		:  str		)  ->				List[str]:


			'''simple docstring'''





			lowerCAmelCase_    :Dict			     =  [v.split("""_"""		)[0].split("""block"""		)[1] for v in original_param_names if v.startswith("""block"""		)]
			lowerCAmelCase_    :str			     =  sorted(set(lowercase__		)		)
			lowerCAmelCase_    :Optional[Any]			     =  len(lowercase__		)
			lowerCAmelCase_    :Any			     =  {b: str(lowercase__		) for b, i in zip(lowercase__      ,     range(lowercase__		)		)}

			lowerCAmelCase_    :List[str]			     =  []
			rename_keys.append(("""stem_conv/kernel:0""", """embeddings.convolution.weight""")		)
			rename_keys.append(("""stem_bn/gamma:0""", """embeddings.batchnorm.weight""")		)
			rename_keys.append(("""stem_bn/beta:0""", """embeddings.batchnorm.bias""")		)
			rename_keys.append(("""stem_bn/moving_mean:0""", """embeddings.batchnorm.running_mean""")		)
			rename_keys.append(("""stem_bn/moving_variance:0""", """embeddings.batchnorm.running_var""")		)

			for b in block_names:
						lowerCAmelCase_    :Dict			     =  block_name_mapping[b]
						rename_keys.append((f"""block{b}_expand_conv/kernel:0""", f"""encoder.blocks.{hf_b}.expansion.expand_conv.weight""")		)
						rename_keys.append((f"""block{b}_expand_bn/gamma:0""", f"""encoder.blocks.{hf_b}.expansion.expand_bn.weight""")		)
						rename_keys.append((f"""block{b}_expand_bn/beta:0""", f"""encoder.blocks.{hf_b}.expansion.expand_bn.bias""")		)
						rename_keys.append(
						    (f"""block{b}_expand_bn/moving_mean:0""", f"""encoder.blocks.{hf_b}.expansion.expand_bn.running_mean""")		)
						rename_keys.append(
						    (f"""block{b}_expand_bn/moving_variance:0""", f"""encoder.blocks.{hf_b}.expansion.expand_bn.running_var""")		)
						rename_keys.append(
						    (f"""block{b}_dwconv/depthwise_kernel:0""", f"""encoder.blocks.{hf_b}.depthwise_conv.depthwise_conv.weight""")		)
						rename_keys.append((f"""block{b}_bn/gamma:0""", f"""encoder.blocks.{hf_b}.depthwise_conv.depthwise_norm.weight""")		)
						rename_keys.append((f"""block{b}_bn/beta:0""", f"""encoder.blocks.{hf_b}.depthwise_conv.depthwise_norm.bias""")		)
						rename_keys.append(
						    (f"""block{b}_bn/moving_mean:0""", f"""encoder.blocks.{hf_b}.depthwise_conv.depthwise_norm.running_mean""")		)
						rename_keys.append(
						    (f"""block{b}_bn/moving_variance:0""", f"""encoder.blocks.{hf_b}.depthwise_conv.depthwise_norm.running_var""")		)

						rename_keys.append((f"""block{b}_se_reduce/kernel:0""", f"""encoder.blocks.{hf_b}.squeeze_excite.reduce.weight""")		)
						rename_keys.append((f"""block{b}_se_reduce/bias:0""", f"""encoder.blocks.{hf_b}.squeeze_excite.reduce.bias""")		)
						rename_keys.append((f"""block{b}_se_expand/kernel:0""", f"""encoder.blocks.{hf_b}.squeeze_excite.expand.weight""")		)
						rename_keys.append((f"""block{b}_se_expand/bias:0""", f"""encoder.blocks.{hf_b}.squeeze_excite.expand.bias""")		)
						rename_keys.append(
						    (f"""block{b}_project_conv/kernel:0""", f"""encoder.blocks.{hf_b}.projection.project_conv.weight""")		)
						rename_keys.append((f"""block{b}_project_bn/gamma:0""", f"""encoder.blocks.{hf_b}.projection.project_bn.weight""")		)
						rename_keys.append((f"""block{b}_project_bn/beta:0""", f"""encoder.blocks.{hf_b}.projection.project_bn.bias""")		)
						rename_keys.append(
						    (f"""block{b}_project_bn/moving_mean:0""", f"""encoder.blocks.{hf_b}.projection.project_bn.running_mean""")		)
						rename_keys.append(
						    (f"""block{b}_project_bn/moving_variance:0""", f"""encoder.blocks.{hf_b}.projection.project_bn.running_var""")		)

			rename_keys.append(("""top_conv/kernel:0""", """encoder.top_conv.weight""")		)
			rename_keys.append(("""top_bn/gamma:0""", """encoder.top_bn.weight""")		)
			rename_keys.append(("""top_bn/beta:0""", """encoder.top_bn.bias""")		)
			rename_keys.append(("""top_bn/moving_mean:0""", """encoder.top_bn.running_mean""")		)
			rename_keys.append(("""top_bn/moving_variance:0""", """encoder.top_bn.running_var""")		)

			lowerCAmelCase_    :str			     =  {}
			for item in rename_keys:
						if item[0] in original_param_names:
									lowerCAmelCase_    :List[str]			     =  """efficientnet.""" + item[1]

			lowerCAmelCase_    :Optional[int]			     =  """classifier.weight"""
			lowerCAmelCase_    :Union[str, Any]			     =  """classifier.bias"""
			return key_mapping






def 						_snake_case       (	lowercase__		:  List[str]      ,     lowercase__		:  List[Any]      ,     lowercase__		:  List[Any]		)  ->				Union[str, Any]:


			'''simple docstring'''





			for key, value in tf_params.items():
						if "normalization" in key:
									continue

						lowerCAmelCase_    :Dict			     =  key_mapping[key]
						if "_conv" in key and "kernel" in key:
									lowerCAmelCase_    :Union[str, Any]			     =  torch.from_numpy(lowercase__		).permute(3      ,     2      ,     0      ,     1		)
						elif "depthwise_kernel" in key:
									lowerCAmelCase_    :List[str]			     =  torch.from_numpy(lowercase__		).permute(2      ,     3      ,     0      ,     1		)
						elif "kernel" in key:
									lowerCAmelCase_    :List[Any]			     =  torch.from_numpy(np.transpose(lowercase__		)		)
						else:
									lowerCAmelCase_    :Any			     =  torch.from_numpy(lowercase__		)

						# Replace HF parameters with original TF model parameters
						assert hf_params[hf_key].shape == new_hf_value.shape
						hf_params[hf_key].copy_(lowercase__		)






@torch.no_grad()
def 						_snake_case       (	lowercase__		:  Optional[Any]      ,     lowercase__		:  Union[str, Any]      ,     lowercase__		:  int      ,     lowercase__		:  Optional[int]		)  ->				Optional[int]:


			'''simple docstring'''





			lowerCAmelCase_    :Dict			     =  model_classes[model_name](
			    include_top=lowercase__      ,     weights="""imagenet"""      ,     input_tensor=lowercase__      ,     input_shape=lowercase__      ,     pooling=lowercase__      ,     classes=1_0_0_0      ,     classifier_activation="""softmax"""      ,     )

			lowerCAmelCase_    :List[str]			     =  original_model.trainable_variables
			lowerCAmelCase_    :str			     =  original_model.non_trainable_variables
			lowerCAmelCase_    :Optional[Any]			     =  {param.name: param.numpy() for param in tf_params}
			for param in tf_non_train_params:
						lowerCAmelCase_    :Dict			     =  param.numpy()
			lowerCAmelCase_    :Union[str, Any]			     =  list(tf_params.keys()		)

			# Load HuggingFace model
			lowerCAmelCase_    :Dict			     =  get_efficientnet_config(lowercase__		)
			lowerCAmelCase_    :Union[str, Any]			     =  EfficientNetForImageClassification(lowercase__		).eval()
			lowerCAmelCase_    :int			     =  hf_model.state_dict()

			# Create src-to-dst parameter name mapping dictionary
			print("""Converting parameters..."""		)
			lowerCAmelCase_    :Optional[int]			     =  rename_keys(lowercase__		)
			replace_params(lowercase__      ,     lowercase__      ,     lowercase__		)

			# Initialize preprocessor and preprocess input image
			lowerCAmelCase_    :Dict			     =  convert_image_processor(lowercase__		)
			lowerCAmelCase_    :str			     =  preprocessor(images=prepare_img()      ,     return_tensors="""pt"""		)

			# HF model inference
			hf_model.eval()
			with torch.no_grad():
						lowerCAmelCase_    :str			     =  hf_model(**lowercase__		)
			lowerCAmelCase_    :Tuple			     =  outputs.logits.detach().numpy()

			# Original model inference
			lowerCAmelCase_    :Tuple			     =  False
			lowerCAmelCase_    :Optional[int]			     =  CONFIG_MAP[model_name]["""image_size"""]
			lowerCAmelCase_    :Optional[Any]			     =  prepare_img().resize((image_size, image_size)      ,     resample=PIL.Image.NEAREST		)
			lowerCAmelCase_    :str			     =  image.img_to_array(lowercase__		)
			lowerCAmelCase_    :List[Any]			     =  np.expand_dims(lowercase__      ,     axis=0		)
			lowerCAmelCase_    :Optional[Any]			     =  original_model.predict(lowercase__		)

			# Check whether original and HF model outputs match  -> np.allclose
			assert np.allclose(lowercase__      ,     lowercase__      ,     atol=1E-3		), "The predicted logits are not the same."
			print("""Model outputs match!"""		)

			if save_model:
						# Create folder to save model
						if not os.path.isdir(lowercase__		):
									os.mkdir(lowercase__		)
						# Save converted model and image processor
						hf_model.save_pretrained(lowercase__		)
						preprocessor.save_pretrained(lowercase__		)

			if push_to_hub:
						# Push model and image processor to hub
						print(f"""Pushing converted {model_name} to the hub..."""		)
						lowerCAmelCase_    :List[Any]			     =  f"""efficientnet-{model_name}"""
						preprocessor.push_to_hub(lowercase__		)
						hf_model.push_to_hub(lowercase__		)


if __name__ == "__main__":
				__UpperCAmelCase       =  argparse.ArgumentParser()
				# Required parameters
				parser.add_argument(
				    '--model_name',
				    default='b0',
				    type=str,
				    help='Version name of the EfficientNet model you want to convert, select from [b0, b1, b2, b3, b4, b5, b6, b7].',
				)
				parser.add_argument(
				    '--pytorch_dump_folder_path',
				    default='hf_model',
				    type=str,
				    help='Path to the output PyTorch model directory.',
				)
				parser.add_argument('--save_model', action='store_true', help='Save model to local')
				parser.add_argument('--push_to_hub', action='store_true', help='Push model and image processor to the hub')

				__UpperCAmelCase       =  parser.parse_args()
				convert_efficientnet_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.save_model, args.push_to_hub)










'''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    __snake_case       (   unittest.TestCase):







    """simple docstring"""



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

        lowerCAmelCase_			:      Dict			=   [t[-1] for t in os.walk(os.path.join(lowerCamelCase    ,						os.listdir(lowerCamelCase		)[0]    ,						"""snapshots"""		)		)]
        lowerCAmelCase_			:      List[Any]			=   [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    __snake_case       (   unittest.TestCase):







    """simple docstring"""



    def __lowercase   (   self    :  Any		)   ->  Optional[int]:
      lowerCAmelCase_,							lowerCAmelCase_			:      Union[str, Any]			=   FlaxStableDiffusionPipeline.from_pretrained(
          """hf-internal-testing/tiny-stable-diffusion-pipe"""    ,						safety_checker=lowerCamelCase		)

      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_			:      Tuple			=   jax.random.PRNGKey(0		)
      lowerCAmelCase_			:      Optional[Any]			=   4

      lowerCAmelCase_			:      Dict			=   jax.device_count()
      lowerCAmelCase_			:      Tuple			=   num_samples * [prompt]
      lowerCAmelCase_			:      Optional[int]			=   pipeline.prepare_inputs(lowerCamelCase		)

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

      lowerCAmelCase_			:      str			=   pipeline(lowerCamelCase    ,						lowerCamelCase    ,						lowerCamelCase    ,						lowerCamelCase    ,						jit=lowerCamelCase		).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(lowerCamelCase    ,						dtype=np.floataa		).sum() - 49_947.875		) < 5E-1

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

    def __lowercase   (   self    :  Optional[int]		)   ->  List[str]:
      lowerCAmelCase_,							lowerCAmelCase_			:      Tuple			=   FlaxStableDiffusionPipeline.from_pretrained(
          """CompVis/stable-diffusion-v1-4"""    ,						revision="""flax"""    ,						safety_checker=lowerCamelCase		)

      lowerCAmelCase_			:      Dict			=   (
          """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_			:      List[Any]			=   50

      lowerCAmelCase_			:      str			=   jax.device_count()
      lowerCAmelCase_			:      Optional[Any]			=   num_samples * [prompt]
      lowerCAmelCase_			:      Dict			=   pipeline.prepare_inputs(lowerCamelCase		)

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

      lowerCAmelCase_			:      Optional[int]			=   pipeline(lowerCamelCase    ,						lowerCamelCase    ,						lowerCamelCase    ,						lowerCamelCase    ,						jit=lowerCamelCase		).images

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

    def __lowercase   (   self    :  List[Any]		)   ->  List[Any]:
      lowerCAmelCase_,							lowerCAmelCase_			:      Optional[Any]			=   FlaxStableDiffusionPipeline.from_pretrained(
          """CompVis/stable-diffusion-v1-4"""    ,						revision="""bf16"""    ,						dtype=jnp.bfloataa    ,						safety_checker=lowerCamelCase		)

      lowerCAmelCase_			:      Optional[int]			=   (
          """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[str]			=   jax.random.PRNGKey(0		)
      lowerCAmelCase_			:      Union[str, Any]			=   50

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

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

      lowerCAmelCase_			:      Any			=   pipeline(lowerCamelCase    ,						lowerCamelCase    ,						lowerCamelCase    ,						lowerCamelCase    ,						jit=lowerCamelCase		).images

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

    def __lowercase   (   self    :  int		)   ->  Optional[int]:
      lowerCAmelCase_,							lowerCAmelCase_			:      int			=   FlaxStableDiffusionPipeline.from_pretrained(
          """CompVis/stable-diffusion-v1-4"""    ,						revision="""bf16"""    ,						dtype=jnp.bfloataa		)

      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_			:      str			=   jax.random.PRNGKey(0		)
      lowerCAmelCase_			:      List[Any]			=   50

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

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

      lowerCAmelCase_			:      Tuple			=   pipeline(lowerCamelCase    ,						lowerCamelCase    ,						lowerCamelCase    ,						lowerCamelCase    ,						jit=lowerCamelCase		).images

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

    def __lowercase   (   self    :  List[str]		)   ->  Any:
      lowerCAmelCase_			:      List[Any]			=   FlaxDDIMScheduler(
          beta_start=0.00_085    ,						beta_end=0.012    ,						beta_schedule="""scaled_linear"""    ,						set_alpha_to_one=lowerCamelCase    ,						steps_offset=1    ,						)

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

      lowerCAmelCase_			:      Optional[Any]			=   scheduler_state

      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_			:      int			=   jax.random.PRNGKey(0		)
      lowerCAmelCase_			:      List[str]			=   50

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

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

      lowerCAmelCase_			:      Tuple			=   pipeline(lowerCamelCase    ,						lowerCamelCase    ,						lowerCamelCase    ,						lowerCamelCase    ,						jit=lowerCamelCase		).images

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





    def __lowercase   (   self    :  Any		)   ->  Union[str, Any]:
      lowerCAmelCase_			:      int			=   (
          """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.device_count()
      lowerCAmelCase_			:      Optional[Any]			=   num_samples * [prompt]
      lowerCAmelCase_			:      Union[str, Any]			=   jax.random.split(jax.random.PRNGKey(0		)    ,						lowerCamelCase		)

      lowerCAmelCase_,							lowerCAmelCase_			:      Union[str, Any]			=   FlaxStableDiffusionPipeline.from_pretrained(
          """CompVis/stable-diffusion-v1-4"""    ,						revision="""bf16"""    ,						dtype=jnp.bfloataa    ,						safety_checker=lowerCamelCase    ,						)

      lowerCAmelCase_			:      Optional[int]			=   replicate(lowerCamelCase		)
      lowerCAmelCase_			:      Optional[int]			=   pipeline.prepare_inputs(lowerCamelCase		)
      lowerCAmelCase_			:      Optional[int]			=   shard(lowerCamelCase		)
      lowerCAmelCase_			:      Tuple			=   pipeline(lowerCamelCase    ,						lowerCamelCase    ,						lowerCamelCase    ,						jit=lowerCamelCase		).images
      assert images.shape == (num_samples, 1, 5_12, 5_12, 3)
      lowerCAmelCase_			:      Dict			=   images[2, 0, 2_56, 10:17, 1]

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

      lowerCAmelCase_			:      List[str]			=   replicate(lowerCamelCase		)
      lowerCAmelCase_			:      Dict			=   pipeline.prepare_inputs(lowerCamelCase		)
      lowerCAmelCase_			:      Union[str, Any]			=   shard(lowerCamelCase		)
      lowerCAmelCase_			:      Optional[int]			=   pipeline(lowerCamelCase    ,						lowerCamelCase    ,						lowerCamelCase    ,						jit=lowerCamelCase		).images
      assert images_eff.shape == (num_samples, 1, 5_12, 5_12, 3)
      lowerCAmelCase_			:      List[Any]			=   images[2, 0, 2_56, 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



