






import argparse
import json

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

from transformers import SegformerImageProcessor, SwinConfig, UperNetConfig, UperNetForSemanticSegmentation


def __A    (      __lowerCAmelCase	)->  Dict:





						"""simple docstring"""



						_UpperCAmelCase									=			384
						_UpperCAmelCase									=			7
						if "tiny" in model_name:
												_UpperCAmelCase									=			96
												_UpperCAmelCase									=			(2, 2, 6, 2)
												_UpperCAmelCase									=			(3, 6, 12, 24)
						elif "small" in model_name:
												_UpperCAmelCase									=			96
												_UpperCAmelCase									=			(2, 2, 18, 2)
												_UpperCAmelCase									=			(3, 6, 12, 24)
						elif "base" in model_name:
												_UpperCAmelCase									=			128
												_UpperCAmelCase									=			(2, 2, 18, 2)
												_UpperCAmelCase									=			(4, 8, 16, 32)
												_UpperCAmelCase									=			12
												_UpperCAmelCase									=			512
						elif "large" in model_name:
												_UpperCAmelCase									=			192
												_UpperCAmelCase									=			(2, 2, 18, 2)
												_UpperCAmelCase									=			(6, 12, 24, 48)
												_UpperCAmelCase									=			12
												_UpperCAmelCase									=			768

						# set label information
						_UpperCAmelCase									=			150
						_UpperCAmelCase									=			'huggingface/label-files'
						_UpperCAmelCase									=			'ade20k-id2label.json'
						_UpperCAmelCase									=			json.load(open(hf_hub_download(__lowerCAmelCase ,		__lowerCAmelCase ,		repo_type='dataset'	) ,		'r'	)	)
						_UpperCAmelCase									=			{int(__lowerCAmelCase	): v for k, v in idalabel.items()}
						_UpperCAmelCase									=			{v: k for k, v in idalabel.items()}

						_UpperCAmelCase									=			SwinConfig(
						    embed_dim=__lowerCAmelCase ,		depths=__lowerCAmelCase ,		num_heads=__lowerCAmelCase ,		window_size=__lowerCAmelCase ,		out_features=['stage1', 'stage2', 'stage3', 'stage4'] ,		)
						_UpperCAmelCase									=			UperNetConfig(
						    backbone_config=__lowerCAmelCase ,		auxiliary_in_channels=__lowerCAmelCase ,		num_labels=__lowerCAmelCase ,		idalabel=__lowerCAmelCase ,		labelaid=__lowerCAmelCase ,		)

						return config


def __A    (      __lowerCAmelCase	)->  Tuple:





						"""simple docstring"""



						_UpperCAmelCase									=			[]

						# fmt: off
						# stem
						rename_keys.append(('backbone.patch_embed.projection.weight', 'backbone.embeddings.patch_embeddings.projection.weight')	)
						rename_keys.append(('backbone.patch_embed.projection.bias', 'backbone.embeddings.patch_embeddings.projection.bias')	)
						rename_keys.append(('backbone.patch_embed.norm.weight', 'backbone.embeddings.norm.weight')	)
						rename_keys.append(('backbone.patch_embed.norm.bias', 'backbone.embeddings.norm.bias')	)
						# stages
						for i in range(len(config.backbone_config.depths	)	):
												for j in range(config.backbone_config.depths[i]	):
																		rename_keys.append((F"""backbone.stages.{i}.blocks.{j}.norm1.weight""", F"""backbone.encoder.layers.{i}.blocks.{j}.layernorm_before.weight""")	)
																		rename_keys.append((F"""backbone.stages.{i}.blocks.{j}.norm1.bias""", F"""backbone.encoder.layers.{i}.blocks.{j}.layernorm_before.bias""")	)
																		rename_keys.append((F"""backbone.stages.{i}.blocks.{j}.attn.w_msa.relative_position_bias_table""", F"""backbone.encoder.layers.{i}.blocks.{j}.attention.self.relative_position_bias_table""")	)
																		rename_keys.append((F"""backbone.stages.{i}.blocks.{j}.attn.w_msa.relative_position_index""", F"""backbone.encoder.layers.{i}.blocks.{j}.attention.self.relative_position_index""")	)
																		rename_keys.append((F"""backbone.stages.{i}.blocks.{j}.attn.w_msa.proj.weight""", F"""backbone.encoder.layers.{i}.blocks.{j}.attention.output.dense.weight""")	)
																		rename_keys.append((F"""backbone.stages.{i}.blocks.{j}.attn.w_msa.proj.bias""", F"""backbone.encoder.layers.{i}.blocks.{j}.attention.output.dense.bias""")	)
																		rename_keys.append((F"""backbone.stages.{i}.blocks.{j}.norm2.weight""", F"""backbone.encoder.layers.{i}.blocks.{j}.layernorm_after.weight""")	)
																		rename_keys.append((F"""backbone.stages.{i}.blocks.{j}.norm2.bias""", F"""backbone.encoder.layers.{i}.blocks.{j}.layernorm_after.bias""")	)
																		rename_keys.append((F"""backbone.stages.{i}.blocks.{j}.ffn.layers.0.0.weight""", F"""backbone.encoder.layers.{i}.blocks.{j}.intermediate.dense.weight""")	)
																		rename_keys.append((F"""backbone.stages.{i}.blocks.{j}.ffn.layers.0.0.bias""", F"""backbone.encoder.layers.{i}.blocks.{j}.intermediate.dense.bias""")	)
																		rename_keys.append((F"""backbone.stages.{i}.blocks.{j}.ffn.layers.1.weight""", F"""backbone.encoder.layers.{i}.blocks.{j}.output.dense.weight""")	)
																		rename_keys.append((F"""backbone.stages.{i}.blocks.{j}.ffn.layers.1.bias""", F"""backbone.encoder.layers.{i}.blocks.{j}.output.dense.bias""")	)

												if i < 3:
																		rename_keys.append((F"""backbone.stages.{i}.downsample.reduction.weight""", F"""backbone.encoder.layers.{i}.downsample.reduction.weight""")	)
																		rename_keys.append((F"""backbone.stages.{i}.downsample.norm.weight""", F"""backbone.encoder.layers.{i}.downsample.norm.weight""")	)
																		rename_keys.append((F"""backbone.stages.{i}.downsample.norm.bias""", F"""backbone.encoder.layers.{i}.downsample.norm.bias""")	)
												rename_keys.append((F"""backbone.norm{i}.weight""", F"""backbone.hidden_states_norms.stage{i+1}.weight""")	)
												rename_keys.append((F"""backbone.norm{i}.bias""", F"""backbone.hidden_states_norms.stage{i+1}.bias""")	)

						# decode head
						rename_keys.extend(
						    [
						        ('decode_head.conv_seg.weight', 'decode_head.classifier.weight'),
						        ('decode_head.conv_seg.bias', 'decode_head.classifier.bias'),
						        ('auxiliary_head.conv_seg.weight', 'auxiliary_head.classifier.weight'),
						        ('auxiliary_head.conv_seg.bias', 'auxiliary_head.classifier.bias'),
						    ]	)
						# fmt: on

						return rename_keys


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





						"""simple docstring"""



						_UpperCAmelCase									=			dct.pop(__lowerCAmelCase	)
						_UpperCAmelCase									=			val


def __A    (      __lowerCAmelCase ,		__lowerCAmelCase	)->  List[Any]:





						"""simple docstring"""



						_UpperCAmelCase									=			[int(backbone_config.embed_dim * 2**i	) for i in range(len(backbone_config.depths	)	)]
						for i in range(len(backbone_config.depths	)	):
												_UpperCAmelCase									=			num_features[i]
												for j in range(backbone_config.depths[i]	):
																		# fmt: off
																		# read in weights + bias of input projection layer (in original implementation, this is a single matrix + bias)
																		_UpperCAmelCase									=			state_dict.pop(F"""backbone.stages.{i}.blocks.{j}.attn.w_msa.qkv.weight"""	)
																		_UpperCAmelCase									=			state_dict.pop(F"""backbone.stages.{i}.blocks.{j}.attn.w_msa.qkv.bias"""	)
																		# next, add query, keys and values (in that order) to the state dict
																		_UpperCAmelCase									=			in_proj_weight[:dim, :]
																		_UpperCAmelCase									=			in_proj_bias[: dim]
																		_UpperCAmelCase									=			in_proj_weight[
																		    dim : dim * 2, :
																		]
																		_UpperCAmelCase									=			in_proj_bias[
																		    dim : dim * 2
																		]
																		_UpperCAmelCase									=			in_proj_weight[
																		    -dim :, :
																		]
																		_UpperCAmelCase									=			in_proj_bias[-dim :]
																		# fmt: on


def __A    (      __lowerCAmelCase	)->  Optional[Any]:





						"""simple docstring"""



						_UpperCAmelCase					, _UpperCAmelCase									=			x.shape
						_UpperCAmelCase									=			x.reshape(__lowerCAmelCase ,		4 ,		in_channel // 4	)
						_UpperCAmelCase									=			x[:, [0, 2, 1, 3], :].transpose(1 ,		2	).reshape(__lowerCAmelCase ,		__lowerCAmelCase	)
						return x


def __A    (      __lowerCAmelCase	)->  Optional[Any]:





						"""simple docstring"""



						_UpperCAmelCase					, _UpperCAmelCase									=			x.shape
						_UpperCAmelCase									=			x.reshape(__lowerCAmelCase ,		in_channel // 4 ,		4	)
						_UpperCAmelCase									=			x[:, :, [0, 2, 1, 3]].transpose(1 ,		2	).reshape(__lowerCAmelCase ,		__lowerCAmelCase	)

						return x


def __A    (      __lowerCAmelCase	)->  Optional[Any]:





						"""simple docstring"""



						_UpperCAmelCase									=			x.shape[0]
						_UpperCAmelCase									=			x.reshape(4 ,		in_channel // 4	)
						_UpperCAmelCase									=			x[[0, 2, 1, 3], :].transpose(0 ,		1	).reshape(__lowerCAmelCase	)
						return x


def __A    (      __lowerCAmelCase	)->  Dict:





						"""simple docstring"""



						_UpperCAmelCase									=			x.shape[0]
						_UpperCAmelCase									=			x.reshape(in_channel // 4 ,		4	)
						_UpperCAmelCase									=			x[:, [0, 2, 1, 3]].transpose(0 ,		1	).reshape(__lowerCAmelCase	)
						return x


def __A    (      __lowerCAmelCase ,		__lowerCAmelCase ,		__lowerCAmelCase	)->  int:





						"""simple docstring"""



						_UpperCAmelCase									=			{
						    'upernet-swin-tiny': 'https://download.openmmlab.com/mmsegmentation/v0.5/swin/upernet_swin_tiny_patch4_window7_512x512_160k_ade20k_pretrain_224x224_1K/upernet_swin_tiny_patch4_window7_512x512_160k_ade20k_pretrain_224x224_1K_20210531_112542-e380ad3e.pth',
						    'upernet-swin-small': 'https://download.openmmlab.com/mmsegmentation/v0.5/swin/upernet_swin_small_patch4_window7_512x512_160k_ade20k_pretrain_224x224_1K/upernet_swin_small_patch4_window7_512x512_160k_ade20k_pretrain_224x224_1K_20210526_192015-ee2fff1c.pth',
						    'upernet-swin-base': 'https://download.openmmlab.com/mmsegmentation/v0.5/swin/upernet_swin_base_patch4_window12_512x512_160k_ade20k_pretrain_384x384_22K/upernet_swin_base_patch4_window12_512x512_160k_ade20k_pretrain_384x384_22K_20210531_125459-429057bf.pth',
						    'upernet-swin-large': 'https://download.openmmlab.com/mmsegmentation/v0.5/swin/upernet_swin_large_patch4_window12_512x512_pretrain_384x384_22K_160k_ade20k/upernet_swin_large_patch4_window12_512x512_pretrain_384x384_22K_160k_ade20k_20220318_091743-9ba68901.pth',
						}
						_UpperCAmelCase									=			model_name_to_url[model_name]
						_UpperCAmelCase									=			torch.hub.load_state_dict_from_url(__lowerCAmelCase ,		map_location='cpu' ,		file_name=__lowerCAmelCase	)[
						    'state_dict'
						]

						for name, param in state_dict.items():
												print(__lowerCAmelCase ,		param.shape	)

						_UpperCAmelCase									=			get_upernet_config(__lowerCAmelCase	)
						_UpperCAmelCase									=			UperNetForSemanticSegmentation(__lowerCAmelCase	)
						model.eval()

						# replace "bn" => "batch_norm"
						for key in state_dict.copy().keys():
												_UpperCAmelCase									=			state_dict.pop(__lowerCAmelCase	)
												if "bn" in key:
																		_UpperCAmelCase									=			key.replace('bn' ,		'batch_norm'	)
												_UpperCAmelCase									=			val

						# rename keys
						_UpperCAmelCase									=			create_rename_keys(__lowerCAmelCase	)
						for src, dest in rename_keys:
												rename_key(__lowerCAmelCase ,		__lowerCAmelCase ,		__lowerCAmelCase	)
						read_in_q_k_v(__lowerCAmelCase ,		config.backbone_config	)

						# fix downsample parameters
						for key, value in state_dict.items():
												if "downsample" in key:
																		if "reduction" in key:
																								_UpperCAmelCase									=			reverse_correct_unfold_reduction_order(__lowerCAmelCase	)
																		if "norm" in key:
																								_UpperCAmelCase									=			reverse_correct_unfold_norm_order(__lowerCAmelCase	)

						model.load_state_dict(__lowerCAmelCase	)

						# verify on image
						_UpperCAmelCase									=			'https://huggingface.co/datasets/hf-internal-testing/fixtures_ade20k/resolve/main/ADE_val_00000001.jpg'
						_UpperCAmelCase									=			Image.open(requests.get(__lowerCAmelCase ,		stream=__lowerCAmelCase	).raw	).convert('RGB'	)

						_UpperCAmelCase									=			SegformerImageProcessor()
						_UpperCAmelCase									=			processor(__lowerCAmelCase ,		return_tensors='pt'	).pixel_values

						with torch.no_grad():
												_UpperCAmelCase									=			model(__lowerCAmelCase	)
												_UpperCAmelCase									=			outputs.logits

						print(logits.shape	)
						print('First values of logits:' ,		logits[0, 0, :3, :3]	)
						# assert values
						if model_name == "upernet-swin-tiny":
												_UpperCAmelCase									=			torch.tensor(
												    [[-7.59_58, -7.59_58, -7.43_02], [-7.59_58, -7.59_58, -7.43_02], [-7.47_97, -7.47_97, -7.30_68]]	)
						elif model_name == "upernet-swin-small":
												_UpperCAmelCase									=			torch.tensor(
												    [[-7.19_21, -7.19_21, -6.95_32], [-7.19_21, -7.19_21, -6.95_32], [-7.09_08, -7.09_08, -6.85_34]]	)
						elif model_name == "upernet-swin-base":
												_UpperCAmelCase									=			torch.tensor(
												    [[-6.58_51, -6.58_51, -6.43_30], [-6.58_51, -6.58_51, -6.43_30], [-6.47_63, -6.47_63, -6.32_54]]	)
						elif model_name == "upernet-swin-large":
												_UpperCAmelCase									=			torch.tensor(
												    [[-7.52_97, -7.52_97, -7.38_02], [-7.52_97, -7.52_97, -7.38_02], [-7.40_44, -7.40_44, -7.25_86]]	)
						print('Logits:' ,		outputs.logits[0, 0, :3, :3]	)
						assert torch.allclose(outputs.logits[0, 0, :3, :3] ,		__lowerCAmelCase ,		atol=1E-4	)
						print('Looks ok!'	)

						if pytorch_dump_folder_path is not None:
												print(F"""Saving model {model_name} to {pytorch_dump_folder_path}"""	)
												model.save_pretrained(__lowerCAmelCase	)
												print(F"""Saving processor to {pytorch_dump_folder_path}"""	)
												processor.save_pretrained(__lowerCAmelCase	)

						if push_to_hub:
												print(F"""Pushing model and processor for {model_name} to hub"""	)
												model.push_to_hub(F"""openmmlab/{model_name}"""	)
												processor.push_to_hub(F"""openmmlab/{model_name}"""	)


if __name__ == "__main__":
		_a      	=  argparse.ArgumentParser()
		# Required parameters
		parser.add_argument(
		    '''--model_name''',
		    default='''upernet-swin-tiny''',
		    type=str,
		    choices=[F'''upernet-swin-{size}''' for size in ['''tiny''', '''small''', '''base''', '''large''']],
		    help='''Name of the Swin + UperNet model you\'d like to convert.''',
		)
		parser.add_argument(
		    '''--pytorch_dump_folder_path''', default=None, type=str, help='''Path to the output PyTorch model directory.'''
		)
		parser.add_argument(
		    '''--push_to_hub''', action='''store_true''', help='''Whether or not to push the converted model to the 🤗 hub.'''
		)

		_a      	=  parser.parse_args()
		convert_upernet_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)








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
