



"""simple docstring"""


import argparse
import os
from pathlib import Path

import torch
from bark.generation import _load_model as _bark_load_model
from huggingface_hub import hf_hub_download

from transformers import EncodecConfig, EncodecModel, set_seed
from transformers.models.bark.configuration_bark import (
    BarkCoarseConfig,
    BarkConfig,
    BarkFineConfig,
    BarkSemanticConfig,
)
from transformers.models.bark.generation_configuration_bark import (
    BarkCoarseGenerationConfig,
    BarkFineGenerationConfig,
    BarkGenerationConfig,
    BarkSemanticGenerationConfig,
)
from transformers.models.bark.modeling_bark import BarkCoarseModel, BarkFineModel, BarkModel, BarkSemanticModel
from transformers.utils import logging


logging.set_verbosity_info()
__A   :							Any										=   logging.get_logger(__name__)

set_seed(770)


__A   :							List[Any]										=   {
    "c_attn": "att_proj",
    "c_proj": "out_proj",
    "c_fc": "in_proj",
    "transformer.": "",
    "h.": "layers.",
    "ln_1": "layernorm_1",
    "ln_2": "layernorm_2",
    "ln_f": "layernorm_final",
    "wpe": "position_embeds_layer",
    "wte": "input_embeds_layer",
}


__A   :							int										=   {
    "text_small": {
        "repo_id": "suno/bark",
        "file_name": "text.pt",
    },
    "coarse_small": {
        "repo_id": "suno/bark",
        "file_name": "coarse.pt",
    },
    "fine_small": {
        "repo_id": "suno/bark",
        "file_name": "fine.pt",
    },
    "text": {
        "repo_id": "suno/bark",
        "file_name": "text_2.pt",
    },
    "coarse": {
        "repo_id": "suno/bark",
        "file_name": "coarse_2.pt",
    },
    "fine": {
        "repo_id": "suno/bark",
        "file_name": "fine_2.pt",
    },
}

__A   :							List[str]										=   os.path.dirname(os.path.abspath(__file__))
__A   :							List[str]										=   os.path.join(os.path.expanduser("~"), ".cache")
__A   :							Union[str, Any]										=   os.path.join(os.getenv("XDG_CACHE_HOME", default_cache_dir), "suno", "bark_v0")







def  lowercase	(      _SCREAMING_SNAKE_CASE   :						Optional[Any]					,     _SCREAMING_SNAKE_CASE   :						Dict=False				):






				'''simple docstring'''



				_UpperCAmelCase					=    model_type
				if use_small:
								key += "_small"
				return os.path.join(_SCREAMING_SNAKE_CASE					,     REMOTE_MODEL_PATHS[key]['''file_name''']				)







def  lowercase	(      _SCREAMING_SNAKE_CASE   :						List[Any]					,     _SCREAMING_SNAKE_CASE   :						Union[str, Any]				):






				'''simple docstring'''



				os.makedirs(_SCREAMING_SNAKE_CASE					,     exist_ok=_SCREAMING_SNAKE_CASE				)
				hf_hub_download(repo_id=_SCREAMING_SNAKE_CASE					,     filename=_SCREAMING_SNAKE_CASE					,     local_dir=_SCREAMING_SNAKE_CASE				)







def  lowercase	(      _SCREAMING_SNAKE_CASE   :						str					,     _SCREAMING_SNAKE_CASE   :						Optional[Any]					,     _SCREAMING_SNAKE_CASE   :						Dict=False					,     _SCREAMING_SNAKE_CASE   :						Union[str, Any]="text"				):






				'''simple docstring'''



				if model_type == "text":
								_UpperCAmelCase					=    BarkSemanticModel
								_UpperCAmelCase					=    BarkSemanticConfig
								_UpperCAmelCase					=    BarkSemanticGenerationConfig
				elif model_type == "coarse":
								_UpperCAmelCase					=    BarkCoarseModel
								_UpperCAmelCase					=    BarkCoarseConfig
								_UpperCAmelCase					=    BarkCoarseGenerationConfig
				elif model_type == "fine":
								_UpperCAmelCase					=    BarkFineModel
								_UpperCAmelCase					=    BarkFineConfig
								_UpperCAmelCase					=    BarkFineGenerationConfig
				else:
								raise NotImplementedError()
				_UpperCAmelCase					=    f'{model_type}_small' if use_small else model_type
				_UpperCAmelCase					=    REMOTE_MODEL_PATHS[model_key]
				if not os.path.exists(_SCREAMING_SNAKE_CASE				):
								logger.info(f'{model_type} model not found, downloading into `{CACHE_DIR}`.'				)
								_download(model_info['''repo_id''']					,     model_info['''file_name''']				)
				_UpperCAmelCase					=    torch.load(_SCREAMING_SNAKE_CASE					,     map_location=_SCREAMING_SNAKE_CASE				)
				# this is a hack
				_UpperCAmelCase					=    checkpoint['''model_args''']
				if "input_vocab_size" not in model_args:
								_UpperCAmelCase					=    model_args['''vocab_size''']
								_UpperCAmelCase					=    model_args['''vocab_size''']
								del model_args["vocab_size"]

				# convert Bark model arguments to HF Bark model arguments
				_UpperCAmelCase					=    model_args.pop('''n_head'''				)
				_UpperCAmelCase					=    model_args.pop('''n_embd'''				)
				_UpperCAmelCase					=    model_args.pop('''n_layer'''				)

				_UpperCAmelCase					=    ConfigClass(**checkpoint['''model_args''']				)
				_UpperCAmelCase					=    ModelClass(config=_SCREAMING_SNAKE_CASE				)
				_UpperCAmelCase					=    GenerationConfigClass()

				_UpperCAmelCase					=    model_generation_config
				_UpperCAmelCase					=    checkpoint['''model''']
				# fixup checkpoint
				_UpperCAmelCase					=    '''_orig_mod.'''
				for k, v in list(state_dict.items()				):
								if k.startswith(_SCREAMING_SNAKE_CASE				):
												# replace part of the key with corresponding layer name in HF implementation
												_UpperCAmelCase					=    k[len(_SCREAMING_SNAKE_CASE				) :]
												for old_layer_name in new_layer_name_dict:
																_UpperCAmelCase					=    new_k.replace(_SCREAMING_SNAKE_CASE					,     new_layer_name_dict[old_layer_name]				)

												_UpperCAmelCase					=    state_dict.pop(_SCREAMING_SNAKE_CASE				)

				_UpperCAmelCase					=    set(state_dict.keys()				) - set(model.state_dict().keys()				)
				_UpperCAmelCase					=    {k for k in extra_keys if not k.endswith('''.attn.bias'''				)}
				_UpperCAmelCase					=    set(model.state_dict().keys()				) - set(state_dict.keys()				)
				_UpperCAmelCase					=    {k for k in missing_keys if not k.endswith('''.attn.bias'''				)}
				if len(_SCREAMING_SNAKE_CASE				) != 0:
								raise ValueError(f'extra keys found: {extra_keys}'				)
				if len(_SCREAMING_SNAKE_CASE				) != 0:
								raise ValueError(f'missing keys: {missing_keys}'				)
				model.load_state_dict(_SCREAMING_SNAKE_CASE					,     strict=_SCREAMING_SNAKE_CASE				)
				_UpperCAmelCase					=    model.num_parameters(exclude_embeddings=_SCREAMING_SNAKE_CASE				)
				_UpperCAmelCase					=    checkpoint['''best_val_loss'''].item()
				logger.info(f'model loaded: {round(n_params/1E6					,     1				)}M params, {round(_SCREAMING_SNAKE_CASE					,     3				)} loss'				)
				model.eval()
				model.to(_SCREAMING_SNAKE_CASE				)
				del checkpoint, state_dict

				return model







def  lowercase	(      _SCREAMING_SNAKE_CASE   :						Optional[int]					,     _SCREAMING_SNAKE_CASE   :						Optional[int]=False					,     _SCREAMING_SNAKE_CASE   :						Union[str, Any]="text"				):






				'''simple docstring'''



				if model_type not in ("text", "coarse", "fine"):
								raise NotImplementedError()

				_UpperCAmelCase					=    '''cpu'''  # do conversion on cpu

				_UpperCAmelCase					=    _get_ckpt_path(_SCREAMING_SNAKE_CASE					,     use_small=_SCREAMING_SNAKE_CASE				)
				_UpperCAmelCase					=    _load_model(_SCREAMING_SNAKE_CASE					,     _SCREAMING_SNAKE_CASE					,     model_type=_SCREAMING_SNAKE_CASE					,     use_small=_SCREAMING_SNAKE_CASE				)

				# load bark initial model
				_UpperCAmelCase					=    _bark_load_model(_SCREAMING_SNAKE_CASE					,     '''cpu'''					,     model_type=_SCREAMING_SNAKE_CASE					,     use_small=_SCREAMING_SNAKE_CASE				)

				if model_type == "text":
								_UpperCAmelCase					=    bark_model['''model''']

				if model.num_parameters(exclude_embeddings=_SCREAMING_SNAKE_CASE				) != bark_model.get_num_params():
								raise ValueError('''initial and new models don\'t have the same number of parameters'''				)

				# check if same output as the bark model
				_UpperCAmelCase					=    5
				_UpperCAmelCase					=    10

				if model_type in ["text", "coarse"]:
								_UpperCAmelCase					=    torch.randint(256					,     (batch_size, sequence_length)					,     dtype=torch.int				)
								_UpperCAmelCase					=    bark_model(_SCREAMING_SNAKE_CASE				)[0]

								_UpperCAmelCase					=    model(_SCREAMING_SNAKE_CASE				)

								# take last logits
								_UpperCAmelCase					=    output_new_model_total.logits[:, [-1], :]

				else:
								_UpperCAmelCase					=    3
								_UpperCAmelCase					=    8
								_UpperCAmelCase					=    torch.randint(256					,     (batch_size, sequence_length, n_codes_total)					,     dtype=torch.int				)

								_UpperCAmelCase					=    model(_SCREAMING_SNAKE_CASE					,     _SCREAMING_SNAKE_CASE				)
								_UpperCAmelCase					=    bark_model(_SCREAMING_SNAKE_CASE					,     _SCREAMING_SNAKE_CASE				)

								_UpperCAmelCase					=    output_new_model_total.logits

				# output difference should come from the difference of self-attention implementation design
				if output_new_model.shape != output_old_model.shape:
								raise ValueError('''initial and new outputs don\'t have the same shape'''				)
				if (output_new_model - output_old_model).abs().max().item() > 1E-3:
								raise ValueError('''initial and new outputs are not equal'''				)

				Path(_SCREAMING_SNAKE_CASE				).mkdir(exist_ok=_SCREAMING_SNAKE_CASE				)
				model.save_pretrained(_SCREAMING_SNAKE_CASE				)







def  lowercase	(      _SCREAMING_SNAKE_CASE   :						str					,     _SCREAMING_SNAKE_CASE   :						Union[str, Any]					,     _SCREAMING_SNAKE_CASE   :						Optional[Any]					,     _SCREAMING_SNAKE_CASE   :						List[str]					,     _SCREAMING_SNAKE_CASE   :						Union[str, Any]					,     _SCREAMING_SNAKE_CASE   :						int					,     ):






				'''simple docstring'''



				_UpperCAmelCase					=    os.path.join(_SCREAMING_SNAKE_CASE					,     _SCREAMING_SNAKE_CASE				)

				_UpperCAmelCase					=    BarkSemanticConfig.from_pretrained(os.path.join(_SCREAMING_SNAKE_CASE					,     '''config.json'''				)				)
				_UpperCAmelCase					=    BarkCoarseConfig.from_pretrained(os.path.join(_SCREAMING_SNAKE_CASE					,     '''config.json'''				)				)
				_UpperCAmelCase					=    BarkFineConfig.from_pretrained(os.path.join(_SCREAMING_SNAKE_CASE					,     '''config.json'''				)				)
				_UpperCAmelCase					=    EncodecConfig.from_pretrained('''facebook/encodec_24khz'''				)

				_UpperCAmelCase					=    BarkSemanticModel.from_pretrained(_SCREAMING_SNAKE_CASE				)
				_UpperCAmelCase					=    BarkCoarseModel.from_pretrained(_SCREAMING_SNAKE_CASE				)
				_UpperCAmelCase					=    BarkFineModel.from_pretrained(_SCREAMING_SNAKE_CASE				)
				_UpperCAmelCase					=    EncodecModel.from_pretrained('''facebook/encodec_24khz'''				)

				_UpperCAmelCase					=    BarkConfig.from_sub_model_configs(
				    _SCREAMING_SNAKE_CASE					,     _SCREAMING_SNAKE_CASE					,     _SCREAMING_SNAKE_CASE					,     _SCREAMING_SNAKE_CASE				)

				_UpperCAmelCase					=    BarkGenerationConfig.from_sub_model_configs(
				    semantic.generation_config					,     coarseAcoustic.generation_config					,     fineAcoustic.generation_config				)

				_UpperCAmelCase					=    BarkModel(_SCREAMING_SNAKE_CASE				)

				_UpperCAmelCase					=    semantic
				_UpperCAmelCase					=    coarseAcoustic
				_UpperCAmelCase					=    fineAcoustic
				_UpperCAmelCase					=    codec

				_UpperCAmelCase					=    bark_generation_config

				Path(_SCREAMING_SNAKE_CASE				).mkdir(exist_ok=_SCREAMING_SNAKE_CASE				)
				bark.save_pretrained(_SCREAMING_SNAKE_CASE					,     repo_id=_SCREAMING_SNAKE_CASE					,     push_to_hub=_SCREAMING_SNAKE_CASE				)


if __name__ == "__main__":
					__A   :							List[str]										=   argparse.ArgumentParser()
					# Required parameters

					parser.add_argument("model_type", type=str, help="text, coarse or fine.")
					parser.add_argument("pytorch_dump_folder_path", default=None, type=str, help="Path to the output PyTorch model.")
					parser.add_argument("--is_small", action="store_true", help="convert the small version instead of the large.")

					__A   :							Tuple										=   parser.parse_args()

					load_model(args.pytorch_dump_folder_path, model_type=args.model_type, use_small=args.is_small)








"""simple docstring"""


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 			_a (					lowerCAmelCase			,							lowerCAmelCase			,							unittest.TestCase):
			"""simple docstring"""
			UpperCamelCase__			      =	StableDiffusionXLImgaImgPipeline
			UpperCamelCase__			      =	TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {"""height""", """width"""}
			UpperCamelCase__			      =	PipelineTesterMixin.required_optional_params - {"""latents"""}
			UpperCamelCase__			      =	TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
			UpperCamelCase__			      =	IMAGE_TO_IMAGE_IMAGE_PARAMS
			UpperCamelCase__			      =	IMAGE_TO_IMAGE_IMAGE_PARAMS



			def  lowercase__						(		self     :    Tuple   )->int:
							torch.manual_seed(0   )
							_UpperCAmelCase					=    UNetaDConditionModel(
							    block_out_channels=(3_2, 6_4)				,   layers_per_block=2				,   sample_size=3_2				,   in_channels=4				,   out_channels=4				,   down_block_types=('''DownBlock2D''', '''CrossAttnDownBlock2D''')				,   up_block_types=('''CrossAttnUpBlock2D''', '''UpBlock2D''')				,   attention_head_dim=(2, 4)				,   use_linear_projection=__UpperCamelCase				,   addition_embed_type='''text_time'''				,   addition_time_embed_dim=8				,   transformer_layers_per_block=(1, 2)				,   projection_class_embeddings_input_dim=8_0				,   cross_attention_dim=6_4				,   )
							_UpperCAmelCase					=    EulerDiscreteScheduler(
							    beta_start=0.0_0_0_8_5				,   beta_end=0.0_1_2				,   steps_offset=1				,   beta_schedule='''scaled_linear'''				,   timestep_spacing='''leading'''				,   )
							torch.manual_seed(0   )
							_UpperCAmelCase					=    AutoencoderKL(
							    block_out_channels=[3_2, 6_4]				,   in_channels=3				,   out_channels=3				,   down_block_types=['''DownEncoderBlock2D''', '''DownEncoderBlock2D''']				,   up_block_types=['''UpDecoderBlock2D''', '''UpDecoderBlock2D''']				,   latent_channels=4				,   sample_size=1_2_8				,   )
							torch.manual_seed(0   )
							_UpperCAmelCase					=    CLIPTextConfig(
							    bos_token_id=0				,   eos_token_id=2				,   hidden_size=3_2				,   intermediate_size=3_7				,   layer_norm_eps=1e-05				,   num_attention_heads=4				,   num_hidden_layers=5				,   pad_token_id=1				,   vocab_size=1_0_0_0				,   hidden_act='''gelu'''				,   projection_dim=3_2				,   )
							_UpperCAmelCase					=    CLIPTextModel(__UpperCamelCase   )
							_UpperCAmelCase					=    CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''				,   local_files_only=__UpperCamelCase   )

							_UpperCAmelCase					=    CLIPTextModelWithProjection(__UpperCamelCase   )
							_UpperCAmelCase					=    CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''				,   local_files_only=__UpperCamelCase   )

							_UpperCAmelCase					=    {
							    '''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  lowercase__						(		self     :    int				,   __UpperCamelCase     :    Optional[int]				,   __UpperCamelCase     :    int=0   )->Any:
							_UpperCAmelCase					=    floats_tensor((1, 3, 3_2, 3_2)				,   rng=random.Random(__UpperCamelCase   )   ).to(__UpperCamelCase   )
							_UpperCAmelCase					=    image / 2 + 0.5
							if str(__UpperCamelCase   ).startswith('''mps'''   ):
											_UpperCAmelCase					=    torch.manual_seed(__UpperCamelCase   )
							else:
											_UpperCAmelCase					=    torch.Generator(device=__UpperCamelCase   ).manual_seed(__UpperCamelCase   )
							_UpperCAmelCase					=    {
							    '''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.7_5,
							}
							return inputs



			def  lowercase__						(		self     :    Any   )->int:
							_UpperCAmelCase					=    '''cpu'''  # ensure determinism for the device-dependent torch.Generator
							_UpperCAmelCase					=    self.get_dummy_components()
							_UpperCAmelCase					=    StableDiffusionXLImgaImgPipeline(**__UpperCamelCase   )
							_UpperCAmelCase					=    sd_pipe.to(__UpperCamelCase   )
							sd_pipe.set_progress_bar_config(disable=__UpperCamelCase   )

							_UpperCAmelCase					=    self.get_dummy_inputs(__UpperCamelCase   )
							_UpperCAmelCase					=    sd_pipe(**__UpperCamelCase   ).images
							_UpperCAmelCase					=    image[0, -3:, -3:, -1]

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

							_UpperCAmelCase					=    np.array([0.4_6_5_6, 0.4_8_4_0, 0.4_4_3_9, 0.6_6_9_8, 0.5_5_7_4, 0.4_5_2_4, 0.5_7_9_9, 0.5_9_4_3, 0.5_1_6_5]   )

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



			def  lowercase__						(		self     :    Optional[Any]   )->Optional[int]:
							super().test_attention_slicing_forward_pass(expected_max_diff=3e-3   )



			def  lowercase__						(		self     :    Optional[Any]   )->List[Any]:
							super().test_inference_batch_single_identical(expected_max_diff=3e-3   )



			def  lowercase__						(		self     :    List[str]   )->str:
							pass



			def  lowercase__						(		self     :    Dict   )->List[str]:
							_UpperCAmelCase					=    self.get_dummy_components()
							_UpperCAmelCase					=    StableDiffusionXLImgaImgPipeline(**__UpperCamelCase   )
							_UpperCAmelCase					=    sd_pipe.to(__UpperCamelCase   )
							_UpperCAmelCase					=    sd_pipe.to(__UpperCamelCase   )
							sd_pipe.set_progress_bar_config(disable=__UpperCamelCase   )

							# forward without prompt embeds
							_UpperCAmelCase					=    self.get_dummy_inputs(__UpperCamelCase   )
							_UpperCAmelCase					=    3 * ['''this is a negative prompt''']
							_UpperCAmelCase					=    negative_prompt
							_UpperCAmelCase					=    3 * [inputs['''prompt''']]

							_UpperCAmelCase					=    sd_pipe(**__UpperCamelCase   )
							_UpperCAmelCase					=    output.images[0, -3:, -3:, -1]

							# forward with prompt embeds
							_UpperCAmelCase					=    self.get_dummy_inputs(__UpperCamelCase   )
							_UpperCAmelCase					=    3 * ['''this is a negative prompt''']
							_UpperCAmelCase					=    3 * [inputs.pop('''prompt'''   )]

							(
							    (
							    _UpperCAmelCase
							)		,						(
							    _UpperCAmelCase
							)		,						(
							    _UpperCAmelCase
							)		,						(
							    _UpperCAmelCase
							)		,						
							)					=    sd_pipe.encode_prompt(__UpperCamelCase				,   negative_prompt=__UpperCamelCase   )

							_UpperCAmelCase					=    sd_pipe(
							    **__UpperCamelCase				,   prompt_embeds=__UpperCamelCase				,   negative_prompt_embeds=__UpperCamelCase				,   pooled_prompt_embeds=__UpperCamelCase				,   negative_pooled_prompt_embeds=__UpperCamelCase				,   )
							_UpperCAmelCase					=    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 			_a (					unittest.TestCase):
			"""simple docstring"""



			def  lowercase__						(		self     :    Union[str, Any]   )->Any:
							super().tearDown()
							gc.collect()
							torch.cuda.empty_cache()



			def  lowercase__						(		self     :    Optional[Any]				,   __UpperCamelCase     :    Union[str, Any]				,   __UpperCamelCase     :    Dict="cpu"				,   __UpperCamelCase     :    int=torch.floataa				,   __UpperCamelCase     :    Optional[Any]=0   )->Union[str, Any]:
							_UpperCAmelCase					=    torch.Generator(device=__UpperCamelCase   ).manual_seed(__UpperCamelCase   )
							_UpperCAmelCase					=    np.random.RandomState(__UpperCamelCase   ).standard_normal((1, 4, 6_4, 6_4)   )
							_UpperCAmelCase					=    torch.from_numpy(__UpperCamelCase   ).to(device=__UpperCamelCase				,   dtype=__UpperCamelCase   )
							_UpperCAmelCase					=    {
							    '''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  lowercase__						(		self     :    Optional[Any]   )->Optional[Any]:
							_UpperCAmelCase					=    DiffusionPipeline.from_pretrained('''stabilityai/stable-diffusion-2-base'''   )
							pipe.to(__UpperCamelCase   )
							pipe.set_progress_bar_config(disable=__UpperCamelCase   )

							_UpperCAmelCase					=    self.get_inputs(__UpperCamelCase   )
							_UpperCAmelCase					=    pipe(**__UpperCamelCase   ).images
							_UpperCAmelCase					=    image[0, -3:, -3:, -1].flatten()

							assert image.shape == (1, 5_1_2, 5_1_2, 3)
							_UpperCAmelCase					=    np.array([0.4_9_4_9_3, 0.4_7_8_9_6, 0.4_0_7_9_8, 0.5_4_2_1_4, 0.5_3_2_1_2, 0.4_8_2_0_2, 0.4_7_6_5_6, 0.4_6_3_2_9, 0.4_8_5_0_6]   )
							assert np.abs(image_slice - expected_slice   ).max() < 7e-3



