






import argparse
import collections

import numpy as np
import torch
from flax import traverse_util
from tax import checkpoints

from transformers import MTaConfig, UMTaEncoderModel, UMTaForConditionalGeneration
from transformers.utils import logging


logging.set_verbosity_info()

def    SCREAMING_SNAKE_CASE_							(							__magic_name__					:  List[str]	,			__magic_name__					:  int	,			__magic_name__					:  int					)		->			Tuple:






	"""simple docstring"""



	return params[f"""{prefix}/{prefix}/relpos_bias/rel_embedding"""][:, i, :]

def    SCREAMING_SNAKE_CASE_							(							__magic_name__					:  Optional[int]	,			__magic_name__					:  Optional[Any]	,			__magic_name__					:  Tuple	,			__magic_name__					:  Any="attention"					)		->			str:






	"""simple docstring"""



	UpperCamelCase     :str          =	np.ascontiguousarray(params[f"""{prefix}/{prefix}/{layer_name}/key/kernel"""][:, i, :, :]					)
	UpperCamelCase     :Optional[Any]          =	k_tmp.reshape(k_tmp.shape[0]	,			k_tmp.shape[1] * k_tmp.shape[2]					)
	UpperCamelCase     :Dict          =	np.ascontiguousarray(params[f"""{prefix}/{prefix}/{layer_name}/out/kernel"""][:, i, :, :]					)
	UpperCamelCase     :List[str]          =	o_tmp.reshape(o_tmp.shape[0] * o_tmp.shape[1]	,			o_tmp.shape[2]					)
	UpperCamelCase     :List[str]          =	np.ascontiguousarray(params[f"""{prefix}/{prefix}/{layer_name}/query/kernel"""][:, i, :, :]					)
	UpperCamelCase     :int          =	q_tmp.reshape(q_tmp.shape[0]	,			q_tmp.shape[1] * q_tmp.shape[2]					)
	UpperCamelCase     :Optional[Any]          =	np.ascontiguousarray(params[f"""{prefix}/{prefix}/{layer_name}/value/kernel"""][:, i, :, :]					)
	UpperCamelCase     :int          =	v_tmp.reshape(v_tmp.shape[0]	,			v_tmp.shape[1] * v_tmp.shape[2]					)
	return k, o, q, v

def    SCREAMING_SNAKE_CASE_							(							__magic_name__					:  Dict	,			__magic_name__					:  Optional[int]	,			__magic_name__					:  Optional[Any]	,			__magic_name__					:  Any=False					)		->			Tuple:






	"""simple docstring"""



	if split_mlp_wi:
		UpperCamelCase     :Optional[Any]          =	params[f"""{prefix}/{prefix}/mlp/wi_0/kernel"""][:, i, :]
		UpperCamelCase     :List[str]          =	params[f"""{prefix}/{prefix}/mlp/wi_1/kernel"""][:, i, :]
		UpperCamelCase     :Dict          =	(wi_a, wi_a)
	else:
		UpperCamelCase     :Union[str, Any]          =	params[f"""{prefix}/{prefix}/mlp/wi/kernel"""][:, i, :]

	UpperCamelCase     :Tuple          =	params[f"""{prefix}/{prefix}/mlp/wo/kernel"""][:, i, :]
	return wi, wo

def    SCREAMING_SNAKE_CASE_							(							__magic_name__					:  Dict	,			__magic_name__					:  Dict	,			__magic_name__					:  List[Any]	,			__magic_name__					:  List[str]					)		->			List[str]:






	"""simple docstring"""



	return params[f"""{prefix}/{prefix}/{layer_name}/scale"""][:, i]

def    SCREAMING_SNAKE_CASE_							(							__magic_name__					:  dict	,			*, __magic_name__					:  int	,			__magic_name__					:  bool	,			__magic_name__					:  bool = False					)		->			Dict:






	"""simple docstring"""



	UpperCamelCase     :Dict          =	traverse_util.flatten_dict(variables["""target"""]					)
	UpperCamelCase     :int          =	{"""/""".join(__magic_name__					): v for k, v in old.items()}

	# v1.1 models have a gated GeLU with wi_0 and wi_1 instead of wi
	UpperCamelCase     :str          =	"""encoder/encoder/mlp/wi_0/kernel""" in old
	print("""Split MLP:"""	,			__magic_name__					)

	UpperCamelCase     :Optional[int]          =	collections.OrderedDict()

	# Shared embeddings.
	UpperCamelCase     :Dict          =	old["""token_embedder/embedding"""]

	# Encoder.
	for i in range(__magic_name__					):
		# Block i, layer 0 (Self Attention).
		UpperCamelCase     :Optional[int]          =	tax_layer_norm_lookup(__magic_name__	,			__magic_name__	,			"""encoder"""	,			"""pre_attention_layer_norm"""					)
		UpperCamelCase     ,      UpperCamelCase     ,      UpperCamelCase     ,      UpperCamelCase     :Optional[Any]          =	tax_attention_lookup(__magic_name__	,			__magic_name__	,			"""encoder"""	,			"""attention"""					)
		UpperCamelCase     :Union[str, Any]          =	layer_norm
		UpperCamelCase     :Tuple          =	k.T
		UpperCamelCase     :Optional[Any]          =	o.T
		UpperCamelCase     :Dict          =	q.T
		UpperCamelCase     :List[str]          =	v.T

		# Block i, layer 1 (MLP).
		UpperCamelCase     :Optional[Any]          =	tax_layer_norm_lookup(__magic_name__	,			__magic_name__	,			"""encoder"""	,			"""pre_mlp_layer_norm"""					)
		UpperCamelCase     ,      UpperCamelCase     :Dict          =	tax_mlp_lookup(__magic_name__	,			__magic_name__	,			"""encoder"""	,			__magic_name__					)
		UpperCamelCase     :List[Any]          =	layer_norm
		if split_mlp_wi:
			UpperCamelCase     :Any          =	wi[0].T
			UpperCamelCase     :Optional[Any]          =	wi[1].T
		else:
			UpperCamelCase     :Optional[int]          =	wi.T
		UpperCamelCase     :Any          =	wo.T
		if scalable_attention:
			# convert the rel_embedding of each layer
			UpperCamelCase     :Optional[Any]          =	tax_relpos_bias_lookup(
			    __magic_name__	,			__magic_name__	,			"""encoder"""					).T

	UpperCamelCase     :Tuple          =	old["""encoder/encoder_norm/scale"""]

	if not scalable_attention:
		UpperCamelCase     :Dict          =	tax_relpos_bias_lookup(
		    __magic_name__	,			0	,			"""encoder"""					).T
		UpperCamelCase     :Tuple          =	tax_relpos_bias_lookup(
		    __magic_name__	,			0	,			"""decoder"""					).T

	if not is_encoder_only:
		# Decoder.
		for i in range(__magic_name__					):
			# Block i, layer 0 (Self Attention).
			UpperCamelCase     :Dict          =	tax_layer_norm_lookup(__magic_name__	,			__magic_name__	,			"""decoder"""	,			"""pre_self_attention_layer_norm"""					)
			UpperCamelCase     ,      UpperCamelCase     ,      UpperCamelCase     ,      UpperCamelCase     :Dict          =	tax_attention_lookup(__magic_name__	,			__magic_name__	,			"""decoder"""	,			"""self_attention"""					)
			UpperCamelCase     :int          =	layer_norm
			UpperCamelCase     :Optional[Any]          =	k.T
			UpperCamelCase     :Tuple          =	o.T
			UpperCamelCase     :Optional[Any]          =	q.T
			UpperCamelCase     :int          =	v.T

			# Block i, layer 1 (Cross Attention).
			UpperCamelCase     :str          =	tax_layer_norm_lookup(__magic_name__	,			__magic_name__	,			"""decoder"""	,			"""pre_cross_attention_layer_norm"""					)
			UpperCamelCase     ,      UpperCamelCase     ,      UpperCamelCase     ,      UpperCamelCase     :Tuple          =	tax_attention_lookup(__magic_name__	,			__magic_name__	,			"""decoder"""	,			"""encoder_decoder_attention"""					)
			UpperCamelCase     :str          =	layer_norm
			UpperCamelCase     :List[str]          =	k.T
			UpperCamelCase     :Optional[int]          =	o.T
			UpperCamelCase     :Union[str, Any]          =	q.T
			UpperCamelCase     :Tuple          =	v.T

			# Block i, layer 2 (MLP).
			UpperCamelCase     :Optional[Any]          =	tax_layer_norm_lookup(__magic_name__	,			__magic_name__	,			"""decoder"""	,			"""pre_mlp_layer_norm"""					)
			UpperCamelCase     ,      UpperCamelCase     :Optional[int]          =	tax_mlp_lookup(__magic_name__	,			__magic_name__	,			"""decoder"""	,			__magic_name__					)
			UpperCamelCase     :str          =	layer_norm
			if split_mlp_wi:
				UpperCamelCase     :Any          =	wi[0].T
				UpperCamelCase     :int          =	wi[1].T
			else:
				UpperCamelCase     :List[str]          =	wi.T
			UpperCamelCase     :Union[str, Any]          =	wo.T

			if scalable_attention:
				# convert the rel_embedding of each layer
				UpperCamelCase     :Union[str, Any]          =	tax_relpos_bias_lookup(__magic_name__	,			__magic_name__	,			"""decoder"""					).T

		UpperCamelCase     :int          =	old["""decoder/decoder_norm/scale"""]

		# LM Head (only in v1.1 checkpoints, in v1.0 embeddings are used instead)
		if "decoder/logits_dense/kernel" in old:
			UpperCamelCase     :Optional[Any]          =	old["""decoder/logits_dense/kernel"""].T

	return new

def    SCREAMING_SNAKE_CASE_							(							__magic_name__					:  Optional[int]	,			__magic_name__					:  bool					)		->			List[Any]:






	"""simple docstring"""



	UpperCamelCase     :Union[str, Any]          =	collections.OrderedDict([(k, torch.from_numpy(v.copy()					)) for (k, v) in converted_params.items()]					)

	# Add what is missing.
	if "encoder.embed_tokens.weight" not in state_dict:
		UpperCamelCase     :int          =	state_dict["""shared.weight"""]

	if not is_encoder_only:
		if "decoder.embed_tokens.weight" not in state_dict:
			UpperCamelCase     :List[str]          =	state_dict["""shared.weight"""]

		if "lm_head.weight" not in state_dict:  # For old 1.0 models.
			print("""Using shared word embeddings as lm_head."""					)
			UpperCamelCase     :Optional[Any]          =	state_dict["""shared.weight"""]

	return state_dict

def    SCREAMING_SNAKE_CASE_							(							__magic_name__					:  List[Any]	,			__magic_name__					:  List[Any]	,			__magic_name__					:  List[Any]	,			__magic_name__					:  int	,			__magic_name__					:  List[str]					)		->			int:






	"""simple docstring"""



	UpperCamelCase     :Any          =	checkpoints.load_tax_checkpoint(__magic_name__					)
	UpperCamelCase     :int          =	convert_tax_to_pytorch(
	    __magic_name__	,			num_layers=config.num_layers	,			is_encoder_only=__magic_name__	,			scalable_attention=__magic_name__					)
	UpperCamelCase     :Any          =	make_state_dict(__magic_name__	,			__magic_name__					)
	model.load_state_dict(__magic_name__	,			strict=__magic_name__					)

def    SCREAMING_SNAKE_CASE_							(							__magic_name__					:  str	,			__magic_name__					:  Dict	,			__magic_name__					:  str	,			__magic_name__					:  bool = False	,			__magic_name__					:  bool = False	,			)		->			Dict:






	"""simple docstring"""



	UpperCamelCase     :Optional[Any]          =	MTaConfig.from_json_file(__magic_name__					)
	print(f"""Building PyTorch model from configuration: {config}"""					)
	# Non-v1.1 checkpoints could also use T5Model, but this works for all.
	# The v1.0 checkpoints will simply have an LM head that is the word embeddings.
	if is_encoder_only:
		UpperCamelCase     :List[str]          =	UMTaEncoderModel(__magic_name__					)
	else:
		UpperCamelCase     :List[Any]          =	UMTaForConditionalGeneration(__magic_name__					)

	# Load weights from tf checkpoint
	load_tax_weights_in_ta(__magic_name__	,			__magic_name__	,			__magic_name__	,			__magic_name__	,			__magic_name__					)

	# Save pytorch-model
	print(f"""Save PyTorch model to {pytorch_dump_path}"""					)
	model.save_pretrained(__magic_name__					)

	# Verify that we can load the checkpoint.
	model.from_pretrained(__magic_name__					)
	print("""Done"""					)


if __name__ == "__main__":
	UpperCAmelCase_ :						str      =      argparse.ArgumentParser(description='''Converts a native T5X checkpoint into a PyTorch checkpoint.''')
	# Required parameters
	parser.add_argument(
	    '''--t5x_checkpoint_path''', default=None, type=str, required=True, help='''Path to the T5X checkpoint.'''
	)
	parser.add_argument(
	    '''--config_file''',
	    default=None,
	    type=str,
	    required=True,
	    help='''The config json file corresponding to the pre-trained T5 model.\nThis specifies the model architecture.''',
	)
	parser.add_argument(
	    '''--pytorch_dump_path''', default=None, type=str, required=True, help='''Path to the output PyTorch model.'''
	)
	parser.add_argument(
	    '''--is_encoder_only''', action='''store_true''', help='''Check if the model is encoder-decoder model''', default=False
	)
	parser.add_argument(
	    '''--scalable_attention''',
	    action='''store_true''',
	    help='''Whether the model uses scaled attention (umt5 model)''',
	    default=False,
	)
	UpperCAmelCase_ :						Tuple      =      parser.parse_args()
	convert_tax_checkpoint_to_pytorch(
	    args.tax_checkpoint_path,
	    args.config_file,
	    args.pytorch_dump_path,
	    args.is_encoder_only,
	    args.scalable_attention,
	)








import gc
import random
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTextModelWithProjection, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DiffusionPipeline,
    EulerDiscreteScheduler,
    StableDiffusionXLImgaImgPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import floats_tensor, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..pipeline_params import (
    IMAGE_TO_IMAGE_IMAGE_PARAMS,
    TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS,
    TEXT_GUIDED_IMAGE_VARIATION_PARAMS,
)
from ..test_pipelines_common import PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()

class     _SCREAMING_SNAKE_CASE       (			_a		,   _a		,   unittest.TestCase      ):
	snake_case__							:     Any     	=  StableDiffusionXLImgaImgPipeline
	snake_case__							:     Tuple     	=  TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {"""height""", """width"""}
	snake_case__							:     Tuple     	=  PipelineTesterMixin.required_optional_params - {"""latents"""}
	snake_case__							:     Any     	=  TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
	snake_case__							:     List[str]     	=  IMAGE_TO_IMAGE_IMAGE_PARAMS
	snake_case__							:     Tuple     	=  IMAGE_TO_IMAGE_IMAGE_PARAMS





	def 						_A (		self				:					int	):
		torch.manual_seed(0	)
		UpperCamelCase     :Any          =	UNetaDConditionModel(
		    block_out_channels=(32, 64) ,  layers_per_block=2 ,  sample_size=32 ,  in_channels=4 ,  out_channels=4 ,  down_block_types=("""DownBlock2D""", """CrossAttnDownBlock2D""") ,  up_block_types=("""CrossAttnUpBlock2D""", """UpBlock2D""") ,  attention_head_dim=(2, 4) ,  use_linear_projection=__lowerCamelCase ,  addition_embed_type="""text_time""" ,  addition_time_embed_dim=8 ,  transformer_layers_per_block=(1, 2) ,  projection_class_embeddings_input_dim=80 ,  cross_attention_dim=64 ,  )
		UpperCamelCase     :Tuple          =	EulerDiscreteScheduler(
		    beta_start=0.00085 ,  beta_end=0.012 ,  steps_offset=1 ,  beta_schedule="""scaled_linear""" ,  timestep_spacing="""leading""" ,  )
		torch.manual_seed(0	)
		UpperCamelCase     :Union[str, Any]          =	AutoencoderKL(
		    block_out_channels=[32, 64] ,  in_channels=3 ,  out_channels=3 ,  down_block_types=["""DownEncoderBlock2D""", """DownEncoderBlock2D"""] ,  up_block_types=["""UpDecoderBlock2D""", """UpDecoderBlock2D"""] ,  latent_channels=4 ,  sample_size=128 ,  )
		torch.manual_seed(0	)
		UpperCamelCase     :Optional[int]          =	CLIPTextConfig(
		    bos_token_id=0 ,  eos_token_id=2 ,  hidden_size=32 ,  intermediate_size=37 ,  layer_norm_eps=1E-05 ,  num_attention_heads=4 ,  num_hidden_layers=5 ,  pad_token_id=1 ,  vocab_size=1_000 ,  hidden_act="""gelu""" ,  projection_dim=32 ,  )
		UpperCamelCase     :Any          =	CLIPTextModel(__lowerCamelCase	)
		UpperCamelCase     :List[Any]          =	CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip""" ,  local_files_only=__lowerCamelCase	)

		UpperCamelCase     :List[Any]          =	CLIPTextModelWithProjection(__lowerCamelCase	)
		UpperCamelCase     :int          =	CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip""" ,  local_files_only=__lowerCamelCase	)

		UpperCamelCase     :Union[str, Any]          =	{
		    """unet""": unet,
		    """scheduler""": scheduler,
		    """vae""": vae,
		    """text_encoder""": text_encoder,
		    """tokenizer""": tokenizer,
		    """text_encoder_2""": text_encoder_a,
		    """tokenizer_2""": tokenizer_a,
		    # "safety_checker": None,
		    # "feature_extractor": None,
		}
		return components





	def 						_A (		self				:					Tuple ,  __lowerCamelCase				:					Any ,  __lowerCamelCase				:					Optional[Any]=0	):
		UpperCamelCase     :Tuple          =	floats_tensor((1, 3, 32, 32) ,  rng=random.Random(__lowerCamelCase	)	).to(__lowerCamelCase	)
		UpperCamelCase     :List[str]          =	image / 2 + 0.5
		if str(__lowerCamelCase	).startswith("""mps"""	):
			UpperCamelCase     :Any          =	torch.manual_seed(__lowerCamelCase	)
		else:
			UpperCamelCase     :List[Any]          =	torch.Generator(device=__lowerCamelCase	).manual_seed(__lowerCamelCase	)
		UpperCamelCase     :str          =	{
		    """prompt""": """A painting of a squirrel eating a burger""",
		    """image""": image,
		    """generator""": generator,
		    """num_inference_steps""": 2,
		    """guidance_scale""": 5.0,
		    """output_type""": """numpy""",
		    """strength""": 0.75,
		}
		return inputs





	def 						_A (		self				:					str	):
		UpperCamelCase     :List[str]          =	"""cpu"""  # ensure determinism for the device-dependent torch.Generator
		UpperCamelCase     :Optional[Any]          =	self.get_dummy_components()
		UpperCamelCase     :List[Any]          =	StableDiffusionXLImgaImgPipeline(**__lowerCamelCase	)
		UpperCamelCase     :Any          =	sd_pipe.to(__lowerCamelCase	)
		sd_pipe.set_progress_bar_config(disable=__lowerCamelCase	)

		UpperCamelCase     :Optional[Any]          =	self.get_dummy_inputs(__lowerCamelCase	)
		UpperCamelCase     :Union[str, Any]          =	sd_pipe(**__lowerCamelCase	).images
		UpperCamelCase     :Union[str, Any]          =	image[0, -3:, -3:, -1]

		assert image.shape == (1, 32, 32, 3)

		UpperCamelCase     :List[Any]          =	np.array([0.4656, 0.4840, 0.4439, 0.6698, 0.5574, 0.4524, 0.5799, 0.5943, 0.5165]	)

		assert np.abs(image_slice.flatten() - expected_slice	).max() < 1E-2





	def 						_A (		self				:					Dict	):
		super().test_attention_slicing_forward_pass(expected_max_diff=3E-3	)





	def 						_A (		self				:					Optional[Any]	):
		super().test_inference_batch_single_identical(expected_max_diff=3E-3	)





	def 						_A (		self				:					Union[str, Any]	):
		pass





	def 						_A (		self				:					Optional[int]	):
		UpperCamelCase     :Union[str, Any]          =	self.get_dummy_components()
		UpperCamelCase     :Dict          =	StableDiffusionXLImgaImgPipeline(**__lowerCamelCase	)
		UpperCamelCase     :List[Any]          =	sd_pipe.to(__lowerCamelCase	)
		UpperCamelCase     :List[str]          =	sd_pipe.to(__lowerCamelCase	)
		sd_pipe.set_progress_bar_config(disable=__lowerCamelCase	)

		# forward without prompt embeds
		UpperCamelCase     :List[Any]          =	self.get_dummy_inputs(__lowerCamelCase	)
		UpperCamelCase     :int          =	3 * ["""this is a negative prompt"""]
		UpperCamelCase     :Union[str, Any]          =	negative_prompt
		UpperCamelCase     :Union[str, Any]          =	3 * [inputs["""prompt"""]]

		UpperCamelCase     :Dict          =	sd_pipe(**__lowerCamelCase	)
		UpperCamelCase     :Union[str, Any]          =	output.images[0, -3:, -3:, -1]

		# forward with prompt embeds
		UpperCamelCase     :Union[str, Any]          =	self.get_dummy_inputs(__lowerCamelCase	)
		UpperCamelCase     :Optional[int]          =	3 * ["""this is a negative prompt"""]
		UpperCamelCase     :Union[str, Any]          =	3 * [inputs.pop("""prompt"""	)]

		(
		    (
		    UpperCamelCase
		)     ,      (
		    UpperCamelCase
		)     ,      (
		    UpperCamelCase
		)     ,      (
		    UpperCamelCase
		)     ,      
		)     :Union[str, Any]          =	sd_pipe.encode_prompt(__lowerCamelCase ,  negative_prompt=__lowerCamelCase	)

		UpperCamelCase     :Dict          =	sd_pipe(
		    **__lowerCamelCase ,  prompt_embeds=__lowerCamelCase ,  negative_prompt_embeds=__lowerCamelCase ,  pooled_prompt_embeds=__lowerCamelCase ,  negative_pooled_prompt_embeds=__lowerCamelCase ,  )
		UpperCamelCase     :Union[str, Any]          =	output.images[0, -3:, -3:, -1]

		# make sure that it's equal
		assert np.abs(image_slice_a.flatten() - image_slice_a.flatten()	).max() < 1E-4

@slow
@require_torch_gpu
class     _SCREAMING_SNAKE_CASE       (			unittest.TestCase      ):





	def 						_A (		self				:					Tuple	):
		super().tearDown()
		gc.collect()
		torch.cuda.empty_cache()





	def 						_A (		self				:					List[Any] ,  __lowerCamelCase				:					Any ,  __lowerCamelCase				:					Dict="cpu" ,  __lowerCamelCase				:					List[Any]=torch.floataa ,  __lowerCamelCase				:					Tuple=0	):
		UpperCamelCase     :Optional[int]          =	torch.Generator(device=__lowerCamelCase	).manual_seed(__lowerCamelCase	)
		UpperCamelCase     :Optional[Any]          =	np.random.RandomState(__lowerCamelCase	).standard_normal((1, 4, 64, 64)	)
		UpperCamelCase     :Dict          =	torch.from_numpy(__lowerCamelCase	).to(device=__lowerCamelCase ,  dtype=__lowerCamelCase	)
		UpperCamelCase     :str          =	{
		    """prompt""": """a photograph of an astronaut riding a horse""",
		    """latents""": latents,
		    """generator""": generator,
		    """num_inference_steps""": 3,
		    """guidance_scale""": 7.5,
		    """output_type""": """numpy""",
		}
		return inputs





	def 						_A (		self				:					Optional[Any]	):
		UpperCamelCase     :Any          =	DiffusionPipeline.from_pretrained("""stabilityai/stable-diffusion-2-base"""	)
		pipe.to(__lowerCamelCase	)
		pipe.set_progress_bar_config(disable=__lowerCamelCase	)

		UpperCamelCase     :Optional[Any]          =	self.get_inputs(__lowerCamelCase	)
		UpperCamelCase     :Optional[int]          =	pipe(**__lowerCamelCase	).images
		UpperCamelCase     :Dict          =	image[0, -3:, -3:, -1].flatten()

		assert image.shape == (1, 512, 512, 3)
		UpperCamelCase     :Union[str, Any]          =	np.array([0.49493, 0.47896, 0.40798, 0.54214, 0.53212, 0.48202, 0.47656, 0.46329, 0.48506]	)
		assert np.abs(image_slice - expected_slice	).max() < 7E-3
