





import argparse
import json
from collections import OrderedDict
from pathlib import Path

import requests
import torch
from huggingface_hub import hf_hub_download
from PIL import Image

from transformers import (
    SegformerConfig,
    SegformerForImageClassification,
    SegformerForSemanticSegmentation,
    SegformerImageProcessor,
)
from transformers.utils import logging


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






def      lowerCamelCase_						(				_UpperCamelCase ,				_UpperCamelCase=False )		->						Union[str, Any]:






				"""simple docstring"""


				snake_case_						:    Tuple				=  OrderedDict()
				for key, value in state_dict.items():
								if encoder_only and not key.startswith('''head''' ):
												snake_case_						:    Union[str, Any]				=  '''segformer.encoder.''' + key
								if key.startswith('''backbone''' ):
												snake_case_						:    Any				=  key.replace('''backbone''' ,				'''segformer.encoder''' )
								if "patch_embed" in key:
												# replace for example patch_embed1 by patch_embeddings.0
												snake_case_						:    str				=  key[key.find('''patch_embed''' ) + len('''patch_embed''' )]
												snake_case_						:    List[str]				=  key.replace(f'''patch_embed{idx}''' ,				f'''patch_embeddings.{int(__lowerCAmelCase )-1}''' )
								if "norm" in key:
												snake_case_						:    Any				=  key.replace('''norm''' ,				'''layer_norm''' )
								if "segformer.encoder.layer_norm" in key:
												# replace for example layer_norm1 by layer_norm.0
												snake_case_						:    int				=  key[key.find('''segformer.encoder.layer_norm''' ) + len('''segformer.encoder.layer_norm''' )]
												snake_case_						:    Dict				=  key.replace(f'''layer_norm{idx}''' ,				f'''layer_norm.{int(__lowerCAmelCase )-1}''' )
								if "layer_norm1" in key:
												snake_case_						:    Dict				=  key.replace('''layer_norm1''' ,				'''layer_norm_1''' )
								if "layer_norm2" in key:
												snake_case_						:    Optional[int]				=  key.replace('''layer_norm2''' ,				'''layer_norm_2''' )
								if "block" in key:
												# replace for example block1 by block.0
												snake_case_						:    List[Any]				=  key[key.find('''block''' ) + len('''block''' )]
												snake_case_						:    Union[str, Any]				=  key.replace(f'''block{idx}''' ,				f'''block.{int(__lowerCAmelCase )-1}''' )
								if "attn.q" in key:
												snake_case_						:    Optional[int]				=  key.replace('''attn.q''' ,				'''attention.self.query''' )
								if "attn.proj" in key:
												snake_case_						:    Union[str, Any]				=  key.replace('''attn.proj''' ,				'''attention.output.dense''' )
								if "attn" in key:
												snake_case_						:    Dict				=  key.replace('''attn''' ,				'''attention.self''' )
								if "fc1" in key:
												snake_case_						:    Optional[int]				=  key.replace('''fc1''' ,				'''dense1''' )
								if "fc2" in key:
												snake_case_						:    Tuple				=  key.replace('''fc2''' ,				'''dense2''' )
								if "linear_pred" in key:
												snake_case_						:    Any				=  key.replace('''linear_pred''' ,				'''classifier''' )
								if "linear_fuse" in key:
												snake_case_						:    str				=  key.replace('''linear_fuse.conv''' ,				'''linear_fuse''' )
												snake_case_						:    Tuple				=  key.replace('''linear_fuse.bn''' ,				'''batch_norm''' )
								if "linear_c" in key:
												# replace for example linear_c4 by linear_c.3
												snake_case_						:    Tuple				=  key[key.find('''linear_c''' ) + len('''linear_c''' )]
												snake_case_						:    List[str]				=  key.replace(f'''linear_c{idx}''' ,				f'''linear_c.{int(__lowerCAmelCase )-1}''' )
								if key.startswith('''head''' ):
												snake_case_						:    List[str]				=  key.replace('''head''' ,				'''classifier''' )
								snake_case_						:    str				=  value

				return new_state_dict






def      lowerCamelCase_						(				_UpperCamelCase ,				_UpperCamelCase )		->						List[str]:






				"""simple docstring"""


				for i in range(config.num_encoder_blocks ):
								for j in range(config.depths[i] ):
												# read in weights + bias of keys and values (which is a single matrix in the original implementation)
												snake_case_						:    Optional[int]				=  state_dict.pop(f'''segformer.encoder.block.{i}.{j}.attention.self.kv.weight''' )
												snake_case_						:    int				=  state_dict.pop(f'''segformer.encoder.block.{i}.{j}.attention.self.kv.bias''' )
												# next, add keys and values (in that order) to the state dict
												snake_case_						:    str				=  kv_weight[
												    : config.hidden_sizes[i], :
												]
												snake_case_						:    List[str]				=  kv_bias[: config.hidden_sizes[i]]
												snake_case_						:    List[str]				=  kv_weight[
												    config.hidden_sizes[i] :, :
												]
												snake_case_						:    int				=  kv_bias[
												    config.hidden_sizes[i] :
												]






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






				"""simple docstring"""


				snake_case_						:    Dict				=  '''http://images.cocodataset.org/val2017/000000039769.jpg'''
				snake_case_						:    Dict				=  Image.open(requests.get(__lowerCAmelCase ,				stream=__lowerCAmelCase ).raw )

				return image






@torch.no_grad()
def      lowerCamelCase_						(				_UpperCamelCase ,				_UpperCamelCase ,				_UpperCamelCase )		->						Tuple:






				"""simple docstring"""


				snake_case_						:    Optional[int]				=  SegformerConfig()
				snake_case_						:    List[Any]				=  False

				# set attributes based on model_name
				snake_case_						:    List[Any]				=  '''huggingface/label-files'''
				if "segformer" in model_name:
								snake_case_						:    Tuple				=  model_name[len('''segformer.''' ) : len('''segformer.''' ) + 2]
								if "ade" in model_name:
												snake_case_						:    str				=  150
												snake_case_						:    Optional[Any]				=  '''ade20k-id2label.json'''
												snake_case_						:    int				=  (1, 150, 128, 128)
								elif "city" in model_name:
												snake_case_						:    Any				=  19
												snake_case_						:    Tuple				=  '''cityscapes-id2label.json'''
												snake_case_						:    List[Any]				=  (1, 19, 128, 128)
								else:
												raise ValueError(f'''Model {model_name} not supported''' )
				elif "mit" in model_name:
								snake_case_						:    List[Any]				=  True
								snake_case_						:    List[str]				=  model_name[4:6]
								snake_case_						:    Dict				=  1_000
								snake_case_						:    Optional[Any]				=  '''imagenet-1k-id2label.json'''
								snake_case_						:    Optional[Any]				=  (1, 1_000)
				else:
								raise ValueError(f'''Model {model_name} not supported''' )

				# set config attributes
				snake_case_						:    Tuple				=  json.load(open(hf_hub_download(__lowerCAmelCase ,				__lowerCAmelCase ,				repo_type='''dataset''' ) ,				'''r''' ) )
				snake_case_						:    Optional[int]				=  {int(__lowerCAmelCase ): v for k, v in idalabel.items()}
				snake_case_						:    List[str]				=  idalabel
				snake_case_						:    List[str]				=  {v: k for k, v in idalabel.items()}
				if size == "b0":
								pass
				elif size == "b1":
								snake_case_						:    str				=  [64, 128, 320, 512]
								snake_case_						:    Optional[Any]				=  256
				elif size == "b2":
								snake_case_						:    List[Any]				=  [64, 128, 320, 512]
								snake_case_						:    Union[str, Any]				=  768
								snake_case_						:    int				=  [3, 4, 6, 3]
				elif size == "b3":
								snake_case_						:    List[str]				=  [64, 128, 320, 512]
								snake_case_						:    Optional[int]				=  768
								snake_case_						:    int				=  [3, 4, 18, 3]
				elif size == "b4":
								snake_case_						:    Dict				=  [64, 128, 320, 512]
								snake_case_						:    List[Any]				=  768
								snake_case_						:    Tuple				=  [3, 8, 27, 3]
				elif size == "b5":
								snake_case_						:    Dict				=  [64, 128, 320, 512]
								snake_case_						:    Tuple				=  768
								snake_case_						:    List[Any]				=  [3, 6, 40, 3]
				else:
								raise ValueError(f'''Size {size} not supported''' )

				# load image processor (only resize + normalize)
				snake_case_						:    Optional[Any]				=  SegformerImageProcessor(
				    image_scale=(512, 512) ,				keep_ratio=__lowerCAmelCase ,				align=__lowerCAmelCase ,				do_random_crop=__lowerCAmelCase )

				# prepare image
				snake_case_						:    Optional[Any]				=  prepare_img()
				snake_case_						:    List[Any]				=  image_processor(images=__lowerCAmelCase ,				return_tensors='''pt''' ).pixel_values

				logger.info(f'''Converting model {model_name}...''' )

				# load original state dict
				if encoder_only:
								snake_case_						:    Any				=  torch.load(__lowerCAmelCase ,				map_location=torch.device('''cpu''' ) )
				else:
								snake_case_						:    int				=  torch.load(__lowerCAmelCase ,				map_location=torch.device('''cpu''' ) )['''state_dict''']

				# rename keys
				snake_case_						:    Tuple				=  rename_keys(__lowerCAmelCase ,				encoder_only=__lowerCAmelCase )
				if not encoder_only:
								del state_dict["decode_head.conv_seg.weight"]
								del state_dict["decode_head.conv_seg.bias"]

				# key and value matrices need special treatment
				read_in_k_v(__lowerCAmelCase ,				__lowerCAmelCase )

				# create HuggingFace model and load state dict
				if encoder_only:
								snake_case_						:    List[Any]				=  False
								snake_case_						:    Union[str, Any]				=  SegformerForImageClassification(__lowerCAmelCase )
				else:
								snake_case_						:    Optional[Any]				=  SegformerForSemanticSegmentation(__lowerCAmelCase )
				model.load_state_dict(__lowerCAmelCase )
				model.eval()

				# forward pass
				snake_case_						:    int				=  model(__lowerCAmelCase )
				snake_case_						:    Any				=  outputs.logits

				# set expected_slice based on model name
				# ADE20k checkpoints
				if model_name == "segformer.b0.512x512.ade.160k":
								snake_case_						:    int				=  torch.tensor(
								    [
								        [[-4.6_310, -5.5_232, -6.2_356], [-5.1_921, -6.1_444, -6.5_996], [-5.4_424, -6.2_790, -6.7_574]],
								        [[-12.1_391, -13.3_122, -13.9_554], [-12.8_732, -13.9_352, -14.3_563], [-12.9_438, -13.8_226, -14.2_513]],
								        [[-12.5_134, -13.4_686, -14.4_915], [-12.8_669, -14.4_343, -14.7_758], [-13.2_523, -14.5_819, -15.0_694]],
								    ] )
				elif model_name == "segformer.b1.512x512.ade.160k":
								snake_case_						:    str				=  torch.tensor(
								    [
								        [[-7.5_820, -8.7_231, -8.3_215], [-8.0_600, -10.3_529, -10.0_304], [-7.5_208, -9.4_103, -9.6_239]],
								        [[-12.6_918, -13.8_994, -13.7_137], [-13.3_196, -15.7_523, -15.4_789], [-12.9_343, -14.8_757, -14.9_689]],
								        [[-11.1_911, -11.9_421, -11.3_243], [-11.3_342, -13.6_839, -13.3_581], [-10.3_909, -12.1_832, -12.4_858]],
								    ] )
				elif model_name == "segformer.b2.512x512.ade.160k":
								snake_case_						:    str				=  torch.tensor(
								    [
								        [[-11.8_173, -14.3_850, -16.3_128], [-14.5_648, -16.5_804, -18.6_568], [-14.7_223, -15.7_387, -18.4_218]],
								        [[-15.7_290, -17.9_171, -19.4_423], [-18.3_105, -19.9_448, -21.4_661], [-17.9_296, -18.6_497, -20.7_910]],
								        [[-15.0_783, -17.0_336, -18.2_789], [-16.8_771, -18.6_870, -20.1_612], [-16.2_454, -17.1_426, -19.5_055]],
								    ] )
				elif model_name == "segformer.b3.512x512.ade.160k":
								snake_case_						:    List[str]				=  torch.tensor(
								    [
								        [[-9.0_878, -10.2_081, -10.1_891], [-9.3_144, -10.7_941, -10.9_843], [-9.2_294, -10.3_855, -10.5_704]],
								        [[-12.2_316, -13.9_068, -13.6_102], [-12.9_161, -14.3_702, -14.3_235], [-12.5_233, -13.7_174, -13.7_932]],
								        [[-14.6_275, -15.2_490, -14.9_727], [-14.3_400, -15.9_687, -16.2_827], [-14.1_484, -15.4_033, -15.8_937]],
								    ] )
				elif model_name == "segformer.b4.512x512.ade.160k":
								snake_case_						:    Dict				=  torch.tensor(
								    [
								        [[-12.3_144, -13.2_447, -14.0_802], [-13.3_614, -14.5_816, -15.6_117], [-13.3_340, -14.4_433, -16.2_219]],
								        [[-19.2_781, -20.4_128, -20.7_506], [-20.6_153, -21.6_566, -22.0_998], [-19.9_800, -21.0_430, -22.1_494]],
								        [[-18.8_739, -19.7_804, -21.1_834], [-20.1_233, -21.6_765, -23.2_944], [-20.0_315, -21.2_641, -23.6_944]],
								    ] )
				elif model_name == "segformer.b5.640x640.ade.160k":
								snake_case_						:    List[Any]				=  torch.tensor(
								    [
								        [[-9.5_524, -12.0_835, -11.7_348], [-10.5_229, -13.6_446, -14.5_662], [-9.5_842, -12.8_851, -13.9_414]],
								        [[-15.3_432, -17.5_323, -17.0_818], [-16.3_330, -18.9_255, -19.2_101], [-15.1_340, -17.7_848, -18.3_971]],
								        [[-12.6_072, -14.9_486, -14.6_631], [-13.7_629, -17.0_907, -17.7_745], [-12.7_899, -16.1_695, -17.1_671]],
								    ] )
				# Cityscapes checkpoints
				elif model_name == "segformer.b0.1024x1024.city.160k":
								snake_case_						:    List[Any]				=  torch.tensor(
								    [
								        [[-11.9_295, -13.4_057, -14.8_106], [-13.3_431, -14.8_179, -15.3_781], [-14.2_836, -15.5_942, -16.1_588]],
								        [[-11.4_906, -12.8_067, -13.6_564], [-13.1_189, -14.0_500, -14.1_543], [-13.8_748, -14.5_136, -14.8_789]],
								        [[0.5_374, 0.1_067, -0.4_742], [0.1_141, -0.2_255, -0.7_099], [-0.3_000, -0.5_924, -1.3_105]],
								    ] )
				elif model_name == "segformer.b0.512x1024.city.160k":
								snake_case_						:    Optional[Any]				=  torch.tensor(
								    [
								        [[-7.8_217, -9.8_767, -10.1_717], [-9.4_438, -10.9_058, -11.4_047], [-9.7_939, -12.3_495, -12.1_079]],
								        [[-7.1_514, -9.5_336, -10.0_860], [-9.7_776, -11.6_822, -11.8_439], [-10.1_411, -12.7_655, -12.8_972]],
								        [[0.3_021, 0.0_805, -0.2_310], [-0.0_328, -0.1_605, -0.2_714], [-0.1_408, -0.5_477, -0.6_976]],
								    ] )
				elif model_name == "segformer.b0.640x1280.city.160k":
								snake_case_						:    Any				=  torch.tensor(
								    [
								        [
								            [-1.1372E01, -1.2787E01, -1.3477E01],
								            [-1.2536E01, -1.4194E01, -1.4409E01],
								            [-1.3217E01, -1.4888E01, -1.5327E01],
								        ],
								        [
								            [-1.4791E01, -1.7122E01, -1.8277E01],
								            [-1.7163E01, -1.9192E01, -1.9533E01],
								            [-1.7897E01, -1.9991E01, -2.0315E01],
								        ],
								        [
								            [7.6723E-01, 4.1921E-01, -7.7878E-02],
								            [4.7772E-01, 9.5557E-03, -2.8082E-01],
								            [3.6032E-01, -2.4826E-01, -5.1168E-01],
								        ],
								    ] )
				elif model_name == "segformer.b0.768x768.city.160k":
								snake_case_						:    Any				=  torch.tensor(
								    [
								        [[-9.4_959, -11.3_087, -11.7_479], [-11.0_025, -12.6_540, -12.3_319], [-11.4_064, -13.0_487, -12.9_905]],
								        [[-9.8_905, -11.3_084, -12.0_854], [-11.1_726, -12.7_698, -12.9_583], [-11.5_985, -13.3_278, -14.1_774]],
								        [[0.2_213, 0.0_192, -0.2_466], [-0.1_731, -0.4_213, -0.4_874], [-0.3_126, -0.6_541, -1.1_389]],
								    ] )
				elif model_name == "segformer.b1.1024x1024.city.160k":
								snake_case_						:    Tuple				=  torch.tensor(
								    [
								        [[-13.5_748, -13.9_111, -12.6_500], [-14.3_500, -15.3_683, -14.2_328], [-14.7_532, -16.0_424, -15.6_087]],
								        [[-17.1_651, -15.8_725, -12.9_653], [-17.2_580, -17.3_718, -14.8_223], [-16.6_058, -16.8_783, -16.7_452]],
								        [[-3.6_456, -3.0_209, -1.4_203], [-3.0_797, -3.1_959, -2.0_000], [-1.8_757, -1.9_217, -1.6_997]],
								    ] )
				elif model_name == "segformer.b2.1024x1024.city.160k":
								snake_case_						:    List[Any]				=  torch.tensor(
								    [
								        [[-16.0_976, -16.4_856, -17.3_962], [-16.6_234, -19.0_342, -19.7_685], [-16.0_900, -18.0_661, -19.1_180]],
								        [[-18.4_750, -18.8_488, -19.5_074], [-19.4_030, -22.1_570, -22.5_977], [-19.1_191, -20.8_486, -22.3_783]],
								        [[-4.5_178, -5.5_037, -6.5_109], [-5.0_884, -7.2_174, -8.0_334], [-4.4_156, -5.8_117, -7.2_970]],
								    ] )
				elif model_name == "segformer.b3.1024x1024.city.160k":
								snake_case_						:    Optional[int]				=  torch.tensor(
								    [
								        [[-14.2_081, -14.4_732, -14.1_977], [-14.5_867, -16.4_423, -16.6_356], [-13.4_441, -14.9_685, -16.8_696]],
								        [[-14.4_576, -14.7_073, -15.0_451], [-15.0_816, -17.6_237, -17.9_873], [-14.4_213, -16.0_199, -18.5_992]],
								        [[-4.7_349, -4.9_588, -5.0_966], [-4.3_210, -6.9_325, -7.2_591], [-3.4_312, -4.7_484, -7.1_917]],
								    ] )
				elif model_name == "segformer.b4.1024x1024.city.160k":
								snake_case_						:    str				=  torch.tensor(
								    [
								        [[-11.7_737, -11.9_526, -11.3_273], [-13.6_692, -14.4_574, -13.8_878], [-13.8_937, -14.6_924, -15.9_345]],
								        [[-14.6_706, -14.5_330, -14.1_306], [-16.1_502, -16.8_180, -16.4_269], [-16.8_338, -17.8_939, -20.1_746]],
								        [[1.0_491, 0.8_289, 1.0_310], [1.1_044, 0.5_219, 0.8_055], [1.0_899, 0.6_926, 0.5_590]],
								    ] )
				elif model_name == "segformer.b5.1024x1024.city.160k":
								snake_case_						:    List[Any]				=  torch.tensor(
								    [
								        [[-12.5_641, -13.4_777, -13.0_684], [-13.9_587, -15.8_983, -16.6_557], [-13.3_109, -15.7_350, -16.3_141]],
								        [[-14.7_074, -15.4_352, -14.5_944], [-16.6_353, -18.1_663, -18.6_120], [-15.1_702, -18.0_329, -18.1_547]],
								        [[-1.7_990, -2.0_951, -1.7_784], [-2.6_397, -3.8_245, -3.9_686], [-1.5_264, -2.8_126, -2.9_316]],
								    ] )
				else:
								snake_case_						:    List[Any]				=  logits.argmax(-1 ).item()
								print('''Predicted class:''' ,				model.config.idalabel[predicted_class_idx] )

				# verify logits
				if not encoder_only:
								assert logits.shape == expected_shape
								assert torch.allclose(logits[0, :3, :3, :3] ,				__lowerCAmelCase ,				atol=1E-2 )

				# finally, save model and image processor
				logger.info(f'''Saving PyTorch model and image processor to {pytorch_dump_folder_path}...''' )
				Path(__lowerCAmelCase ).mkdir(exist_ok=__lowerCAmelCase )
				model.save_pretrained(__lowerCAmelCase )
				image_processor.save_pretrained(__lowerCAmelCase )


if __name__ == "__main__":
	lowerCAmelCase_    =				argparse.ArgumentParser()

	parser.add_argument(
	    '''--model_name''',
	    default='''segformer.b0.512x512.ade.160k''',
	    type=str,
	    help='''Name of the model you\'d like to convert.''',
	)
	parser.add_argument(
	    '''--checkpoint_path''', default=None, type=str, help='''Path to the original PyTorch checkpoint (.pth file).'''
	)
	parser.add_argument(
	    '''--pytorch_dump_folder_path''', default=None, type=str, help='''Path to the folder to output PyTorch model.'''
	)
	lowerCAmelCase_    =				parser.parse_args()
	convert_segformer_checkpoint(args.model_name, args.checkpoint_path, args.pytorch_dump_folder_path)








import os
import tempfile
import unittest

import numpy as np

from diffusers.utils import is_flax_available
from diffusers.utils.testing_utils import require_flax, slow


if is_flax_available():
		import jax
		import jax.numpy as jnp
		from flax.jax_utils import replicate
		from flax.training.common_utils import shard

		from diffusers import FlaxDDIMScheduler, FlaxDiffusionPipeline, FlaxStableDiffusionPipeline





@require_flax
class __lowerCamelCase (      unittest.TestCase):





	"""simple docstring"""






	def        UpperCamelCase    (							self    ):
							"""simple docstring"""







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

													_UpperCAmelCase									=			[t[-1] for t in os.walk(os.path.join(UpperCAmelCase       ,    os.listdir(UpperCAmelCase    )[0]       ,    'snapshots'    )    )]
													_UpperCAmelCase									=			[item for sublist in all_root_files for item in sublist]

													# None of the downloaded files should be a PyTorch file even if we have some here:
													# https://huggingface.co/hf-internal-testing/tiny-stable-diffusion-pipe/blob/main/unet/diffusion_pytorch_model.bin
													assert not any(f.endswith('.bin'    ) for f in files    )





@slow
@require_flax
class __lowerCamelCase (      unittest.TestCase):





	"""simple docstring"""






	def        UpperCamelCase    (							self    ):
							"""simple docstring"""







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

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

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

							_UpperCAmelCase									=			jax.device_count()
							_UpperCAmelCase									=			num_samples * [prompt]
							_UpperCAmelCase									=			pipeline.prepare_inputs(UpperCAmelCase    )

							# shard inputs and rng
							_UpperCAmelCase									=			replicate(UpperCAmelCase    )
							_UpperCAmelCase									=			jax.random.split(UpperCAmelCase       ,    UpperCAmelCase    )
							_UpperCAmelCase									=			shard(UpperCAmelCase    )

							_UpperCAmelCase									=			pipeline(UpperCAmelCase       ,    UpperCAmelCase       ,    UpperCAmelCase       ,    UpperCAmelCase       ,    jit=UpperCAmelCase    ).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_51_47_45    ) < 1e-3
													assert np.abs(np.abs(UpperCAmelCase       ,    dtype=np.floataa    ).sum() - 4_99_47.8_75    ) < 5e-1

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

	def        UpperCamelCase    (							self    ):
							"""simple docstring"""







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

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

							_UpperCAmelCase									=			jax.random.PRNGKey(0    )
							_UpperCAmelCase									=			50

							_UpperCAmelCase									=			jax.device_count()
							_UpperCAmelCase									=			num_samples * [prompt]
							_UpperCAmelCase									=			pipeline.prepare_inputs(UpperCAmelCase    )

							# shard inputs and rng
							_UpperCAmelCase									=			replicate(UpperCAmelCase    )
							_UpperCAmelCase									=			jax.random.split(UpperCAmelCase       ,    UpperCAmelCase    )
							_UpperCAmelCase									=			shard(UpperCAmelCase    )

							_UpperCAmelCase									=			pipeline(UpperCAmelCase       ,    UpperCAmelCase       ,    UpperCAmelCase       ,    UpperCAmelCase       ,    jit=UpperCAmelCase    ).images

							assert images.shape == (num_samples, 1, 512, 512, 3)
							if jax.device_count() == 8:
													assert np.abs((np.abs(images[0, 0, :2, :2, -2:]       ,    dtype=np.floataa    ).sum() - 0.05_65_24_01)    ) < 1e-3
													assert np.abs((np.abs(UpperCAmelCase       ,    dtype=np.floataa    ).sum() - 2_38_38_08.2)    ) < 5e-1

	def        UpperCamelCase    (							self    ):
							"""simple docstring"""







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

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

							_UpperCAmelCase									=			jax.random.PRNGKey(0    )
							_UpperCAmelCase									=			50

							_UpperCAmelCase									=			jax.device_count()
							_UpperCAmelCase									=			num_samples * [prompt]
							_UpperCAmelCase									=			pipeline.prepare_inputs(UpperCAmelCase    )

							# shard inputs and rng
							_UpperCAmelCase									=			replicate(UpperCAmelCase    )
							_UpperCAmelCase									=			jax.random.split(UpperCAmelCase       ,    UpperCAmelCase    )
							_UpperCAmelCase									=			shard(UpperCAmelCase    )

							_UpperCAmelCase									=			pipeline(UpperCAmelCase       ,    UpperCAmelCase       ,    UpperCAmelCase       ,    UpperCAmelCase       ,    jit=UpperCAmelCase    ).images

							assert images.shape == (num_samples, 1, 512, 512, 3)
							if jax.device_count() == 8:
													assert np.abs((np.abs(images[0, 0, :2, :2, -2:]       ,    dtype=np.floataa    ).sum() - 0.04_00_39_06)    ) < 1e-3
													assert np.abs((np.abs(UpperCAmelCase       ,    dtype=np.floataa    ).sum() - 2_37_35_16.75)    ) < 5e-1

	def        UpperCamelCase    (							self    ):
							"""simple docstring"""







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

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

							_UpperCAmelCase									=			jax.random.PRNGKey(0    )
							_UpperCAmelCase									=			50

							_UpperCAmelCase									=			jax.device_count()
							_UpperCAmelCase									=			num_samples * [prompt]
							_UpperCAmelCase									=			pipeline.prepare_inputs(UpperCAmelCase    )

							# shard inputs and rng
							_UpperCAmelCase									=			replicate(UpperCAmelCase    )
							_UpperCAmelCase									=			jax.random.split(UpperCAmelCase       ,    UpperCAmelCase    )
							_UpperCAmelCase									=			shard(UpperCAmelCase    )

							_UpperCAmelCase									=			pipeline(UpperCAmelCase       ,    UpperCAmelCase       ,    UpperCAmelCase       ,    UpperCAmelCase       ,    jit=UpperCAmelCase    ).images

							assert images.shape == (num_samples, 1, 512, 512, 3)
							if jax.device_count() == 8:
													assert np.abs((np.abs(images[0, 0, :2, :2, -2:]       ,    dtype=np.floataa    ).sum() - 0.04_00_39_06)    ) < 1e-3
													assert np.abs((np.abs(UpperCAmelCase       ,    dtype=np.floataa    ).sum() - 2_37_35_16.75)    ) < 5e-1

	def        UpperCamelCase    (							self    ):
							"""simple docstring"""







							_UpperCAmelCase									=			FlaxDDIMScheduler(
							    beta_start=0.0_00_85       ,    beta_end=0.0_12       ,    beta_schedule='scaled_linear'       ,    set_alpha_to_one=UpperCAmelCase       ,    steps_offset=1       ,    )

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

							_UpperCAmelCase									=			scheduler_state

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

							_UpperCAmelCase									=			jax.random.PRNGKey(0    )
							_UpperCAmelCase									=			50

							_UpperCAmelCase									=			jax.device_count()
							_UpperCAmelCase									=			num_samples * [prompt]
							_UpperCAmelCase									=			pipeline.prepare_inputs(UpperCAmelCase    )

							# shard inputs and rng
							_UpperCAmelCase									=			replicate(UpperCAmelCase    )
							_UpperCAmelCase									=			jax.random.split(UpperCAmelCase       ,    UpperCAmelCase    )
							_UpperCAmelCase									=			shard(UpperCAmelCase    )

							_UpperCAmelCase									=			pipeline(UpperCAmelCase       ,    UpperCAmelCase       ,    UpperCAmelCase       ,    UpperCAmelCase       ,    jit=UpperCAmelCase    ).images

							assert images.shape == (num_samples, 1, 512, 512, 3)
							if jax.device_count() == 8:
													assert np.abs((np.abs(images[0, 0, :2, :2, -2:]       ,    dtype=np.floataa    ).sum() - 0.0_45_04_39_45)    ) < 1e-3
													assert np.abs((np.abs(UpperCAmelCase       ,    dtype=np.floataa    ).sum() - 2_34_76_93.5)    ) < 5e-1

	def        UpperCamelCase    (							self    ):
							"""simple docstring"""







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

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

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

							_UpperCAmelCase									=			replicate(UpperCAmelCase    )
							_UpperCAmelCase									=			pipeline.prepare_inputs(UpperCAmelCase    )
							_UpperCAmelCase									=			shard(UpperCAmelCase    )
							_UpperCAmelCase									=			pipeline(UpperCAmelCase       ,    UpperCAmelCase       ,    UpperCAmelCase       ,    jit=UpperCAmelCase    ).images
							assert images.shape == (num_samples, 1, 512, 512, 3)
							_UpperCAmelCase									=			images[2, 0, 256, 10:17, 1]

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

							_UpperCAmelCase									=			replicate(UpperCAmelCase    )
							_UpperCAmelCase									=			pipeline.prepare_inputs(UpperCAmelCase    )
							_UpperCAmelCase									=			shard(UpperCAmelCase    )
							_UpperCAmelCase									=			pipeline(UpperCAmelCase       ,    UpperCAmelCase       ,    UpperCAmelCase       ,    jit=UpperCAmelCase    ).images
							assert images_eff.shape == (num_samples, 1, 512, 512, 3)
							_UpperCAmelCase									=			images[2, 0, 256, 10:17, 1]

							# I checked the results visually and they are very similar. However, I saw that the max diff is `1` and the `sum`
							# over the 8 images is exactly `256`, which is very suspicious. Testing a random slice for now.
							assert abs(slice_eff - slice    ).max() < 1e-2
