
"""simple docstring"""





import argparse

import fairseq
import torch
from torch import nn

from transformers import (
    MBartaaTokenizer,
    MBartConfig,
    MBartForCausalLM,
    SpeechEncoderDecoderConfig,
    SpeechEncoderDecoderModel,
    WavaVecaConfig,
    WavaVecaFeatureExtractor,
    WavaVecaModel,
    logging,
)


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

_a														=		{
    """post_extract_proj""": """feature_projection.projection""",
    """encoder.pos_conv.0""": """encoder.pos_conv_embed.conv""",
    """self_attn.k_proj""": """encoder.layers.*.attention.k_proj""",
    """self_attn.v_proj""": """encoder.layers.*.attention.v_proj""",
    """self_attn.q_proj""": """encoder.layers.*.attention.q_proj""",
    """self_attn.out_proj""": """encoder.layers.*.attention.out_proj""",
    """self_attn_layer_norm""": """encoder.layers.*.layer_norm""",
    """fc1""": """encoder.layers.*.feed_forward.intermediate_dense""",
    """fc2""": """encoder.layers.*.feed_forward.output_dense""",
    """final_layer_norm""": """encoder.layers.*.final_layer_norm""",
    """encoder.layer_norm""": """encoder.layer_norm""",
    """w2v_model.layer_norm""": """feature_projection.layer_norm""",
    """quantizer.weight_proj""": """quantizer.weight_proj""",
    """quantizer.vars""": """quantizer.codevectors""",
    """project_q""": """project_q""",
    """final_proj""": """project_hid""",
    """w2v_encoder.proj""": """lm_head""",
    """mask_emb""": """masked_spec_embed""",
}
_a														=		[
    """lm_head""",
    """quantizer.weight_proj""",
    """quantizer.codevectors""",
    """project_q""",
    """project_hid""",
]



def       lowerCamelCase__		(			__snake_case,     __snake_case,     __snake_case,     __snake_case,     __snake_case	)							->		int:


					"""simple docstring"""





					for attribute in key.split('''.'''	):
										_UpperCamelCase							 =  getattr(__snake_case,     __snake_case	)

					if weight_type is not None:
										_UpperCamelCase							 =  getattr(__snake_case,     __snake_case	).shape
					else:
										_UpperCamelCase							 =  hf_pointer.shape

					assert hf_shape == value.shape, (
					    F'''Shape of hf {key + "." + weight_type if weight_type is not None else ""} is {hf_shape}, but should be'''
					    F''' {value.shape} for {full_name}'''
					)

					if weight_type == "weight":
										_UpperCamelCase							 =  value
					elif weight_type == "weight_g":
										_UpperCamelCase							 =  value
					elif weight_type == "weight_v":
										_UpperCamelCase							 =  value
					elif weight_type == "bias":
										_UpperCamelCase							 =  value
					else:
										_UpperCamelCase							 =  value

					logger.info(F'''{key + "." + weight_type if weight_type is not None else ""} was initialized from {full_name}.'''	)



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


					"""simple docstring"""





					_UpperCamelCase							 =  []
					_UpperCamelCase							 =  fairseq_model.state_dict()

					_UpperCamelCase							 =  hf_model.feature_extractor
					_UpperCamelCase							 =  hf_model.adapter

					for name, value in fairseq_dict.items():
										_UpperCamelCase							 =  False
										if "conv_layers" in name:
															load_conv_layer(
															    __snake_case,     __snake_case,     __snake_case,     __snake_case,     hf_model.config.feat_extract_norm == '''group''',     )
															_UpperCamelCase							 =  True
										elif any(x in name for x in ['''adaptor''', '''w2v_encoder.proj.''', '''w2v_proj_ln.''']	):
															load_adapter(__snake_case,     __snake_case,     __snake_case,     __snake_case	)
															_UpperCamelCase							 =  True
										else:
															for key, mapped_key in MAPPING.items():
																				if key in name or key.split('''w2v_model.'''	)[-1] == name.split('''.'''	)[0]:
																									_UpperCamelCase							 =  True
																									if "*" in mapped_key:
																														_UpperCamelCase							 =  name.split(__snake_case	)[0].split('''.'''	)[-2]
																														_UpperCamelCase							 =  mapped_key.replace('''*''',     __snake_case	)
																									if "weight_g" in name:
																														_UpperCamelCase							 =  '''weight_g'''
																									elif "weight_v" in name:
																														_UpperCamelCase							 =  '''weight_v'''
																									elif "bias" in name:
																														_UpperCamelCase							 =  '''bias'''
																									elif "weight" in name:
																														_UpperCamelCase							 =  '''weight'''
																									else:
																														_UpperCamelCase							 =  None
																									set_recursively(__snake_case,     __snake_case,     __snake_case,     __snake_case,     __snake_case	)
																				continue
										if not is_used:
															unused_weights.append(__snake_case	)

					logger.warning(F'''Unused weights: {unused_weights}'''	)



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


					"""simple docstring"""





					_UpperCamelCase							 =  full_name.split('''conv_layers.'''	)[-1]
					_UpperCamelCase							 =  name.split('''.'''	)
					_UpperCamelCase							 =  int(items[0]	)
					_UpperCamelCase							 =  int(items[1]	)

					if type_id == 0:
										if "bias" in name:
															assert value.shape == feature_extractor.conv_layers[layer_id].conv.bias.data.shape, (
															    F'''{full_name} has size {value.shape}, but'''
															    F''' {feature_extractor.conv_layers[layer_id].conv.bias.data.shape} was found.'''
															)
															_UpperCamelCase							 =  value
															logger.info(F'''Feat extract conv layer {layer_id} was initialized from {full_name}.'''	)
										elif "weight" in name:
															assert value.shape == feature_extractor.conv_layers[layer_id].conv.weight.data.shape, (
															    F'''{full_name} has size {value.shape}, but'''
															    F''' {feature_extractor.conv_layers[layer_id].conv.weight.data.shape} was found.'''
															)
															_UpperCamelCase							 =  value
															logger.info(F'''Feat extract conv layer {layer_id} was initialized from {full_name}.'''	)
					elif (type_id == 2 and not use_group_norm) or (type_id == 2 and layer_id == 0 and use_group_norm):
										if "bias" in name:
															assert value.shape == feature_extractor.conv_layers[layer_id].layer_norm.bias.data.shape, (
															    F'''{full_name} has size {value.shape}, but {feature_extractor[layer_id].layer_norm.bias.data.shape} was'''
															    " found."
															)
															_UpperCamelCase							 =  value
															logger.info(F'''Feat extract layer norm weight of layer {layer_id} was initialized from {full_name}.'''	)
										elif "weight" in name:
															assert value.shape == feature_extractor.conv_layers[layer_id].layer_norm.weight.data.shape, (
															    F'''{full_name} has size {value.shape}, but'''
															    F''' {feature_extractor[layer_id].layer_norm.weight.data.shape} was found.'''
															)
															_UpperCamelCase							 =  value
															logger.info(F'''Feat extract layer norm weight of layer {layer_id} was initialized from {full_name}.'''	)
					else:
										unused_weights.append(__snake_case	)



def       lowerCamelCase__		(			__snake_case,     __snake_case,     __snake_case,     __snake_case	)							->		List[str]:


					"""simple docstring"""





					_UpperCamelCase							 =  full_name.split('''adaptor.'''	)[-1]
					_UpperCamelCase							 =  name.split('''.'''	)

					if items[1].isdigit():
										_UpperCamelCase							 =  int(items[1]	)
					else:
										_UpperCamelCase							 =  None

					if "adaptor" not in full_name:
										if "proj_ln" in full_name:
															# has to be layer norm
															if "bias" in name:
																				assert (
																				    value.shape == adapter.proj_layer_norm.bias.data.shape
																				), F'''{full_name} has size {value.shape}, but {adapter.proj_layer_norm.bias.data.shape} was found.'''
																				_UpperCamelCase							 =  value
																				logger.info(F'''Adapter proj layer norm bias was initialized from {full_name}.'''	)
															if "weight" in name:
																				assert (
																				    value.shape == adapter.proj_layer_norm.weight.data.shape
																				), F'''{full_name} has size {value.shape}, but {adapter.proj_layer_norm.weight.data.shape} was found.'''
																				_UpperCamelCase							 =  value
										else:
															# has to be projection layer
															if "bias" in name:
																				assert (
																				    value.shape == adapter.proj.bias.data.shape
																				), F'''{full_name} has size {value.shape}, but {adapter.proj.bias.data.shape} was found.'''
																				_UpperCamelCase							 =  value
																				logger.info(F'''Adapter proj layer bias was initialized from {full_name}.'''	)
															if "weight" in name:
																				assert (
																				    value.shape == adapter.proj.weight.data.shape
																				), F'''{full_name} has size {value.shape}, but {adapter.proj.weight.data.shape} was found.'''
																				_UpperCamelCase							 =  value
																				logger.info(F'''Adapter proj layer weight was initialized from {full_name}.'''	)
					elif isinstance(__snake_case,     __snake_case	):
										if "bias" in name:
															assert (
															    value.shape == adapter.layers[layer_id].conv.bias.data.shape
															), F'''{full_name} has size {value.shape}, but {adapter.layers[layer_id].conv.bias.data.shape} was found.'''
															_UpperCamelCase							 =  value
															logger.info(F'''Adapter layer {layer_id} bias was initialized from {full_name}.'''	)
										elif "weight" in name:
															assert (
															    value.shape == adapter.layers[layer_id].conv.weight.data.shape
															), F'''{full_name} has size {value.shape}, but {adapter.layers[layer_id].conv.weight.data.shape} was found.'''
															_UpperCamelCase							 =  value
															logger.info(F'''Adapter layer {layer_id} bias was initialized from {full_name}.'''	)
					else:
										unused_weights.append(__snake_case	)



def       lowerCamelCase__		(			__snake_case	)							->		Any:


					"""simple docstring"""





					_UpperCamelCase ,  _UpperCamelCase							 =  emb.weight.shape
					_UpperCamelCase							 =  nn.Linear(__snake_case,     __snake_case,     bias=__snake_case	)
					_UpperCamelCase							 =  emb.weight.data
					return lin_layer



@torch.no_grad()
def       lowerCamelCase__		(			__snake_case,     __snake_case,     __snake_case,     __snake_case,     __snake_case,     __snake_case,     __snake_case,     __snake_case,     __snake_case,     __snake_case,     __snake_case,     )							->		Any:


					"""simple docstring"""





					_UpperCamelCase							 =  WavaVecaConfig.from_pretrained(
					    __snake_case,     add_adapter=__snake_case,     adapter_stride=__snake_case,     adapter_kernel_size=__snake_case,     use_auth_token=__snake_case,     output_hidden_size=__snake_case,     )
					_UpperCamelCase							 =  MBartConfig.from_pretrained(__snake_case	)

					# load model
					_UpperCamelCase ,  _UpperCamelCase ,  _UpperCamelCase							 =  fairseq.checkpoint_utils.load_model_ensemble_and_task(
					    [checkpoint_path],     arg_overrides={
					        '''config_yaml''': config_yaml_path,
					        '''data''': '''/'''.join(dict_path.split('''/'''	)[:-1]	),
					        '''w2v_path''': checkpoint_path,
					        '''load_pretrained_decoder_from''': None,
					    },     )
					_UpperCamelCase							 =  model[0].eval()

					# load feature extractor
					_UpperCamelCase							 =  WavaVecaFeatureExtractor.from_pretrained(__snake_case,     use_auth_token=__snake_case	)

					# set weights for wav2vec2 encoder
					_UpperCamelCase							 =  WavaVecaModel(__snake_case	)

					recursively_load_weights_wavaveca(model.encoder,     __snake_case	)

					# load decoder weights
					_UpperCamelCase							 =  MBartForCausalLM(__snake_case	)
					_UpperCamelCase ,  _UpperCamelCase							 =  hf_decoder.model.decoder.load_state_dict(model.decoder.state_dict(),     strict=__snake_case	)
					logger.warning(F'''The following keys are missing when loading the decoder weights: {missing_keys}'''	)
					logger.warning(F'''The following keys are unexpected when loading the decoder weights: {unexpected_keys}'''	)

					_UpperCamelCase							 =  SpeechEncoderDecoderModel(encoder=__snake_case,     decoder=__snake_case	)
					_UpperCamelCase							 =  False

					_UpperCamelCase							 =  MBartaaTokenizer(__snake_case	)
					tokenizer.save_pretrained(__snake_case	)

					_UpperCamelCase							 =  hf_wavavec.config.to_dict()
					_UpperCamelCase							 =  tokenizer.pad_token_id
					_UpperCamelCase							 =  tokenizer.bos_token_id
					_UpperCamelCase							 =  tokenizer.eos_token_id
					_UpperCamelCase							 =  '''mbart50'''
					_UpperCamelCase							 =  '''wav2vec2'''

					_UpperCamelCase							 =  tokenizer.eos_token_id
					_UpperCamelCase							 =  25_00_04
					_UpperCamelCase							 =  tokenizer.eos_token_id

					_UpperCamelCase							 =  SpeechEncoderDecoderConfig.from_dict(__snake_case	)

					hf_wavavec.save_pretrained(__snake_case	)
					feature_extractor.save_pretrained(__snake_case	)


if __name__ == "__main__":
				_a														=		argparse.ArgumentParser()
				parser.add_argument("""--pytorch_dump_folder_path""", default=None, type=str, help="""Path to the output PyTorch model.""")
				parser.add_argument("""--checkpoint_path""", default=None, type=str, help="""Path to fairseq checkpoint""")
				parser.add_argument("""--dict_path""", default=None, type=str, help="""Path to dict of fine-tuned model""")
				parser.add_argument("""--config_yaml_path""", default=None, type=str, help="""Path to yaml file of fine-tuned model""")
				parser.add_argument(
				    """--encoder_config_path""",
				    default="""facebook/wav2vec2-xls-r-1b""",
				    type=str,
				    help="""Path to hf encoder wav2vec2 checkpoint config""",
				)
				parser.add_argument(
				    """--decoder_config_path""",
				    default="""facebook/mbart-large-50-one-to-many-mmt""",
				    type=str,
				    help="""Path to hf decoder checkpoint config""",
				)
				parser.add_argument("""--add_adapter""", default=True, type=bool, help="""whethere to add model adapter layers""")
				parser.add_argument("""--adapter_stride""", default=2, type=int, help="""stride of adapter layers""")
				parser.add_argument("""--adapter_kernel_size""", default=3, type=int, help="""kernel size of adapter layers""")
				parser.add_argument("""--encoder_output_dim""", default=1024, type=int, help="""encoder output dim""")
				parser.add_argument("""--start_token_id""", default=25_0004, type=int, help="""`decoder_start_token_id` of model config""")

				_a														=		parser.parse_args()
				convert_wavaveca_checkpoint(
				    args.checkpoint_path,
				    args.pytorch_dump_folder_path,
				    args.dict_path,
				    args.config_yaml_path,
				    encoder_config_path=args.encoder_config_path,
				    decoder_config_path=args.decoder_config_path,
				    add_adapter=args.add_adapter,
				    adapter_kernel_size=args.adapter_kernel_size,
				    adapter_stride=args.adapter_stride,
				    decoder_start_token_id=args.start_token_id,
				    encoder_output_dim=args.encoder_output_dim,
				)








"""simple docstring"""





import gc
import unittest

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

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    EulerAncestralDiscreteScheduler,
    LMSDiscreteScheduler,
    PNDMScheduler,
    StableDiffusionPanoramaPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu, skip_mps

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()






@skip_mps
class     _UpperCAmelCase(    lowerCamelCase      ,     lowerCamelCase      ,     unittest.TestCase	):
		lowercase__						=      StableDiffusionPanoramaPipeline
		lowercase__						=      TEXT_TO_IMAGE_PARAMS
		lowercase__						=      TEXT_TO_IMAGE_BATCH_PARAMS
		lowercase__						=      TEXT_TO_IMAGE_IMAGE_PARAMS
		lowercase__						=      TEXT_TO_IMAGE_IMAGE_PARAMS



		def 						UpperCAmelCase							(					self)       ->       Dict:






							'''simple docstring'''
							torch.manual_seed(0)
							_UpperCamelCase							 =  UNetaDConditionModel(
							    block_out_channels=(32, 64)				,		layers_per_block=1				,		sample_size=32				,		in_channels=4				,		out_channels=4				,		down_block_types=('''DownBlock2D''', '''CrossAttnDownBlock2D''')				,		up_block_types=('''CrossAttnUpBlock2D''', '''UpBlock2D''')				,		cross_attention_dim=32				,		)
							_UpperCamelCase							 =  DDIMScheduler()
							torch.manual_seed(0)
							_UpperCamelCase							 =  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				,		)
							torch.manual_seed(0)
							_UpperCamelCase							 =  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=10_00				,		)
							_UpperCamelCase							 =  CLIPTextModel(__a)
							_UpperCamelCase							 =  CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip''')

							_UpperCamelCase							 =  {
							    '''unet''': unet,
							    '''scheduler''': scheduler,
							    '''vae''': vae,
							    '''text_encoder''': text_encoder,
							    '''tokenizer''': tokenizer,
							    '''safety_checker''': None,
							    '''feature_extractor''': None,
							}
							return components



		def 						UpperCAmelCase							(					self				,		__a				,		__a=0)       ->       int:






							'''simple docstring'''
							_UpperCamelCase							 =  torch.manual_seed(__a)
							_UpperCamelCase							 =  {
							    '''prompt''': '''a photo of the dolomites''',
							    '''generator''': generator,
							    # Setting height and width to None to prevent OOMs on CPU.
							    '''height''': None,
							    '''width''': None,
							    '''num_inference_steps''': 1,
							    '''guidance_scale''': 6.0,
							    '''output_type''': '''numpy''',
							}
							return inputs



		def 						UpperCAmelCase							(					self)       ->       Dict:






							'''simple docstring'''
							_UpperCamelCase							 =  '''cpu'''  # ensure determinism for the device-dependent torch.Generator
							_UpperCamelCase							 =  self.get_dummy_components()
							_UpperCamelCase							 =  StableDiffusionPanoramaPipeline(**__a)
							_UpperCamelCase							 =  sd_pipe.to(__a)
							sd_pipe.set_progress_bar_config(disable=__a)

							_UpperCamelCase							 =  self.get_dummy_inputs(__a)
							_UpperCamelCase							 =  sd_pipe(**__a).images
							_UpperCamelCase							 =  image[0, -3:, -3:, -1]
							assert image.shape == (1, 64, 64, 3)

							_UpperCamelCase							 =  np.array([0.6186, 0.5374, 0.4915, 0.4135, 0.4114, 0.4563, 0.5128, 0.4977, 0.4757])

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



		def 						UpperCAmelCase							(					self)       ->       Optional[int]:






							'''simple docstring'''
							super().test_inference_batch_consistent(batch_sizes=[1, 2])



		def 						UpperCAmelCase							(					self)       ->       Tuple:






							'''simple docstring'''
							super().test_inference_batch_single_identical(batch_size=2				,		expected_max_diff=3.25e-3)



		def 						UpperCAmelCase							(					self)       ->       List[str]:






							'''simple docstring'''
							_UpperCamelCase							 =  '''cpu'''  # ensure determinism for the device-dependent torch.Generator
							_UpperCamelCase							 =  self.get_dummy_components()
							_UpperCamelCase							 =  StableDiffusionPanoramaPipeline(**__a)
							_UpperCamelCase							 =  sd_pipe.to(__a)
							sd_pipe.set_progress_bar_config(disable=__a)

							_UpperCamelCase							 =  self.get_dummy_inputs(__a)
							_UpperCamelCase							 =  '''french fries'''
							_UpperCamelCase							 =  sd_pipe(**__a				,		negative_prompt=__a)
							_UpperCamelCase							 =  output.images
							_UpperCamelCase							 =  image[0, -3:, -3:, -1]

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

							_UpperCamelCase							 =  np.array([0.6187, 0.5375, 0.4915, 0.4136, 0.4114, 0.4563, 0.5128, 0.4976, 0.4757])

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



		def 						UpperCAmelCase							(					self)       ->       Any:






							'''simple docstring'''
							_UpperCamelCase							 =  '''cpu'''  # ensure determinism for the device-dependent torch.Generator
							_UpperCamelCase							 =  self.get_dummy_components()
							_UpperCamelCase							 =  StableDiffusionPanoramaPipeline(**__a)
							_UpperCamelCase							 =  sd_pipe.to(__a)
							sd_pipe.set_progress_bar_config(disable=__a)

							_UpperCamelCase							 =  self.get_dummy_inputs(__a)
							_UpperCamelCase							 =  sd_pipe(**__a				,		view_batch_size=2)
							_UpperCamelCase							 =  output.images
							_UpperCamelCase							 =  image[0, -3:, -3:, -1]

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

							_UpperCamelCase							 =  np.array([0.6187, 0.5375, 0.4915, 0.4136, 0.4114, 0.4563, 0.5128, 0.4976, 0.4757])

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



		def 						UpperCAmelCase							(					self)       ->       str:






							'''simple docstring'''
							_UpperCamelCase							 =  '''cpu'''  # ensure determinism for the device-dependent torch.Generator
							_UpperCamelCase							 =  self.get_dummy_components()
							_UpperCamelCase							 =  EulerAncestralDiscreteScheduler(
							    beta_start=0.0_0085				,		beta_end=0.012				,		beta_schedule='''scaled_linear''')
							_UpperCamelCase							 =  StableDiffusionPanoramaPipeline(**__a)
							_UpperCamelCase							 =  sd_pipe.to(__a)
							sd_pipe.set_progress_bar_config(disable=__a)

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

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

							_UpperCamelCase							 =  np.array([0.4024, 0.6510, 0.4901, 0.5378, 0.5813, 0.5622, 0.4795, 0.4467, 0.4952])

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




		def 						UpperCAmelCase							(					self)       ->       Any:






							'''simple docstring'''
							_UpperCamelCase							 =  '''cpu'''  # ensure determinism for the device-dependent torch.Generator
							_UpperCamelCase							 =  self.get_dummy_components()
							_UpperCamelCase							 =  PNDMScheduler(
							    beta_start=0.0_0085				,		beta_end=0.012				,		beta_schedule='''scaled_linear'''				,		skip_prk_steps=__a)
							_UpperCamelCase							 =  StableDiffusionPanoramaPipeline(**__a)
							_UpperCamelCase							 =  sd_pipe.to(__a)
							sd_pipe.set_progress_bar_config(disable=__a)

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

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

							_UpperCamelCase							 =  np.array([0.6391, 0.6291, 0.4861, 0.5134, 0.5552, 0.4578, 0.5032, 0.5023, 0.4539])

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






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



		def 						UpperCAmelCase							(					self)       ->       Tuple:






							'''simple docstring'''
							super().tearDown()
							gc.collect()
							torch.cuda.empty_cache()



		def 						UpperCAmelCase							(					self				,		__a=0)       ->       List[str]:






							'''simple docstring'''
							_UpperCamelCase							 =  torch.manual_seed(__a)
							_UpperCamelCase							 =  {
							    '''prompt''': '''a photo of the dolomites''',
							    '''generator''': generator,
							    '''num_inference_steps''': 3,
							    '''guidance_scale''': 7.5,
							    '''output_type''': '''numpy''',
							}
							return inputs



		def 						UpperCAmelCase							(					self)       ->       Any:






							'''simple docstring'''
							_UpperCamelCase							 =  '''stabilityai/stable-diffusion-2-base'''
							_UpperCamelCase							 =  DDIMScheduler.from_pretrained(__a				,		subfolder='''scheduler''')
							_UpperCamelCase							 =  StableDiffusionPanoramaPipeline.from_pretrained(__a				,		scheduler=__a				,		safety_checker=__a)
							pipe.to(__a)
							pipe.set_progress_bar_config(disable=__a)
							pipe.enable_attention_slicing()

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

							assert image.shape == (1, 5_12, 20_48, 3)

							_UpperCamelCase							 =  np.array(
							    [
							        0.3696_8392,
							        0.2702_5372,
							        0.3244_6766,
							        0.2837_9387,
							        0.3636_3274,
							        0.3073_3347,
							        0.2710_0027,
							        0.2705_4125,
							        0.2553_6096,
							    ])

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



		def 						UpperCAmelCase							(					self)       ->       str:






							'''simple docstring'''
							_UpperCamelCase							 =  StableDiffusionPanoramaPipeline.from_pretrained(
							    '''stabilityai/stable-diffusion-2-base'''				,		safety_checker=__a)
							_UpperCamelCase							 =  LMSDiscreteScheduler.from_config(pipe.scheduler.config)
							pipe.to(__a)
							pipe.set_progress_bar_config(disable=__a)
							pipe.enable_attention_slicing()

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

							assert image.shape == (1, 5_12, 20_48, 3)

							_UpperCamelCase							 =  np.array(
							    [
							        [
							            0.0,
							            0.0,
							            0.0,
							            0.0,
							            0.0,
							            0.0,
							            0.0,
							            0.0,
							            0.0,
							        ]
							    ])

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



		def 						UpperCAmelCase							(					self)       ->       Dict:






							'''simple docstring'''
							_UpperCamelCase							 =  0

							def callback_fn(__a				,		__a				,		__a) -> None:
												_UpperCamelCase							 =  True
												nonlocal number_of_steps
												number_of_steps += 1
												if step == 1:
																	_UpperCamelCase							 =  latents.detach().cpu().numpy()
																	assert latents.shape == (1, 4, 64, 2_56)
																	_UpperCamelCase							 =  latents[0, -3:, -3:, -1]

																	_UpperCamelCase							 =  np.array(
																	    [
																	        0.1868_1869,
																	        0.3390_7816,
																	        0.536_1276,
																	        0.1443_2865,
																	        -0.0285_6611,
																	        -0.7394_1123,
																	        0.2339_7987,
																	        0.4732_2682,
																	        -0.3782_3164,
																	    ])
																	assert np.abs(latents_slice.flatten() - expected_slice).max() < 5e-2
												elif step == 2:
																	_UpperCamelCase							 =  latents.detach().cpu().numpy()
																	assert latents.shape == (1, 4, 64, 2_56)
																	_UpperCamelCase							 =  latents[0, -3:, -3:, -1]

																	_UpperCamelCase							 =  np.array(
																	    [
																	        0.1853_9645,
																	        0.3398_7248,
																	        0.537_8559,
																	        0.1443_7142,
																	        -0.0245_5261,
																	        -0.733_8317,
																	        0.2399_0755,
																	        0.4735_6272,
																	        -0.378_6505,
																	    ])

																	assert np.abs(latents_slice.flatten() - expected_slice).max() < 5e-2

							_UpperCamelCase							 =  False

							_UpperCamelCase							 =  '''stabilityai/stable-diffusion-2-base'''
							_UpperCamelCase							 =  DDIMScheduler.from_pretrained(__a				,		subfolder='''scheduler''')
							_UpperCamelCase							 =  StableDiffusionPanoramaPipeline.from_pretrained(__a				,		scheduler=__a				,		safety_checker=__a)
							_UpperCamelCase							 =  pipe.to(__a)
							pipe.set_progress_bar_config(disable=__a)
							pipe.enable_attention_slicing()

							_UpperCamelCase							 =  self.get_inputs()
							pipe(**__a				,		callback=__a				,		callback_steps=1)
							assert callback_fn.has_been_called
							assert number_of_steps == 3




		def 						UpperCAmelCase							(					self)       ->       Optional[int]:






							'''simple docstring'''
							torch.cuda.empty_cache()
							torch.cuda.reset_max_memory_allocated()
							torch.cuda.reset_peak_memory_stats()

							_UpperCamelCase							 =  '''stabilityai/stable-diffusion-2-base'''
							_UpperCamelCase							 =  DDIMScheduler.from_pretrained(__a				,		subfolder='''scheduler''')
							_UpperCamelCase							 =  StableDiffusionPanoramaPipeline.from_pretrained(__a				,		scheduler=__a				,		safety_checker=__a)
							_UpperCamelCase							 =  pipe.to(__a)
							pipe.set_progress_bar_config(disable=__a)
							pipe.enable_attention_slicing(1)
							pipe.enable_sequential_cpu_offload()

							_UpperCamelCase							 =  self.get_inputs()
							_UpperCamelCase							 =  pipe(**__a)

							_UpperCamelCase							 =  torch.cuda.max_memory_allocated()
							# make sure that less than 5.2 GB is allocated
							assert mem_bytes < 5.5 * 10**9






