



"""simple docstring"""







import argparse
import json
import os

import fairseq
import torch
from fairseq.data import Dictionary

from transformers import (
    WavaVecaConformerConfig,
    WavaVecaConformerForCTC,
    WavaVecaConformerForPreTraining,
    WavaVecaCTCTokenizer,
    WavaVecaFeatureExtractor,
    WavaVecaProcessor,
    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.linear_k': 'encoder.layers.*.self_attn.linear_k',
    'self_attn.linear_v': 'encoder.layers.*.self_attn.linear_v',
    'self_attn.linear_q': 'encoder.layers.*.self_attn.linear_q',
    'self_attn.pos_bias_u': 'encoder.layers.*.self_attn.pos_bias_u',
    'self_attn.pos_bias_v': 'encoder.layers.*.self_attn.pos_bias_v',
    'self_attn.linear_out': 'encoder.layers.*.self_attn.linear_out',
    'self_attn.linear_pos': 'encoder.layers.*.self_attn.linear_pos',
    'self_attn.rotary_emb': 'encoder.embed_positions',
    'self_attn_layer_norm': 'encoder.layers.*.self_attn_layer_norm',
    'conv_module.pointwise_conv1': 'encoder.layers.*.conv_module.pointwise_conv1',
    'conv_module.pointwise_conv2': 'encoder.layers.*.conv_module.pointwise_conv2',
    'conv_module.depthwise_conv': 'encoder.layers.*.conv_module.depthwise_conv',
    'conv_module.batch_norm': 'encoder.layers.*.conv_module.batch_norm',
    'conv_module.layer_norm': 'encoder.layers.*.conv_module.layer_norm',
    'ffn1.w_1': 'encoder.layers.*.ffn1.intermediate_dense',
    'ffn1.w_2': 'encoder.layers.*.ffn1.output_dense',
    'ffn1.layer_norm': 'encoder.layers.*.ffn1_layer_norm',
    'ffn2.w_1': 'encoder.layers.*.ffn2.intermediate_dense',
    'ffn2.w_2': 'encoder.layers.*.ffn2.output_dense',
    'ffn2.layer_norm': 'encoder.layers.*.ffn2_layer_norm',
    '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  __SCREAMING_SNAKE_CASE		(			lowerCamelCase_:							Union[str, Any]     ,		lowerCamelCase_:							Tuple     ,		lowerCamelCase_:							Optional[int]     ,		lowerCamelCase_:							Tuple     ,		lowerCamelCase_:							str ):
						"""simple docstring"""






						for attribute in key.split("." ):
												snake_case   :			List[Any]						    = getattr(lowerCamelCase_     ,		lowerCamelCase_ )

						if weight_type is not None:
												snake_case   :			Optional[int]						    = getattr(lowerCamelCase_     ,		lowerCamelCase_ ).shape
						else:
												snake_case   :			Tuple						    = hf_pointer.shape

						if hf_shape != value.shape:
												raise ValueError(
												    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":
												snake_case   :			int						    = value
						elif weight_type == "weight_g":
												snake_case   :			str						    = value
						elif weight_type == "weight_v":
												snake_case   :			Any						    = value
						elif weight_type == "bias":
												snake_case   :			List[Any]						    = value
						elif weight_type == "running_mean":
												snake_case   :			Optional[int]						    = value
						elif weight_type == "running_var":
												snake_case   :			List[Any]						    = value
						elif weight_type == "num_batches_tracked":
												snake_case   :			Union[str, Any]						    = value
						elif weight_type == "inv_freq":
												snake_case   :			Tuple						    = value
						else:
												snake_case   :			str						    = value

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

def  __SCREAMING_SNAKE_CASE		(			lowerCamelCase_:							int     ,		lowerCamelCase_:							Optional[int]     ,		lowerCamelCase_:							Optional[Any] ):
						"""simple docstring"""






						snake_case   :			Tuple						    = []
						snake_case   :			List[str]						    = fairseq_model.state_dict()

						snake_case   :			Union[str, Any]						    = hf_model.wavaveca_conformer.feature_extractor

						for name, value in fairseq_dict.items():
												snake_case   :			List[str]						    = False
												if "conv_layers" in name:
																		load_conv_layer(
																		    lowerCamelCase_     ,		lowerCamelCase_     ,		lowerCamelCase_     ,		lowerCamelCase_     ,		hf_model.config.feat_extract_norm == "group"     ,		)
																		snake_case   :			Dict						    = True
												else:
																		for key, mapped_key in MAPPING.items():
																								snake_case   :			List[str]						    = "wav2vec2_conformer." + mapped_key if mapped_key not in TOP_LEVEL_KEYS else mapped_key
																								if key in name or key.split("w2v_model." )[-1] == name.split("." )[0]:
																														snake_case   :			Dict						    = True
																														if "*" in mapped_key:
																																				snake_case   :			Tuple						    = name.split(lowerCamelCase_ )[0].split("." )[-2]
																																				snake_case   :			int						    = mapped_key.replace("*"     ,		lowerCamelCase_ )
																														if "pos_bias_u" in name:
																																				snake_case   :			Optional[int]						    = None
																														elif "pos_bias_v" in name:
																																				snake_case   :			Dict						    = None
																														elif "weight_g" in name:
																																				snake_case   :			Any						    = "weight_g"
																														elif "weight_v" in name:
																																				snake_case   :			List[str]						    = "weight_v"
																														elif "bias" in name:
																																				snake_case   :			Tuple						    = "bias"
																														elif "weight" in name:
																																				# TODO: don't match quantizer.weight_proj
																																				snake_case   :			Union[str, Any]						    = "weight"
																														elif "running_mean" in name:
																																				snake_case   :			int						    = "running_mean"
																														elif "inv_freq" in name:
																																				snake_case   :			str						    = "inv_freq"
																														elif "running_var" in name:
																																				snake_case   :			Optional[Any]						    = "running_var"
																														elif "num_batches_tracked" in name:
																																				snake_case   :			Tuple						    = "num_batches_tracked"
																														else:
																																				snake_case   :			Optional[int]						    = None
																														set_recursively(lowerCamelCase_     ,		lowerCamelCase_     ,		lowerCamelCase_     ,		lowerCamelCase_     ,		lowerCamelCase_ )
																								continue
												if not is_used:
																		unused_weights.append(lowerCamelCase_ )

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

def  __SCREAMING_SNAKE_CASE		(			lowerCamelCase_:							Optional[Any]     ,		lowerCamelCase_:							Tuple     ,		lowerCamelCase_:							int     ,		lowerCamelCase_:							List[Any]     ,		lowerCamelCase_:							Union[str, Any] ):
						"""simple docstring"""






						snake_case   :			List[Any]						    = full_name.split("conv_layers." )[-1]
						snake_case   :			List[Any]						    = name.split("." )
						snake_case   :			Union[str, Any]						    = int(items[0] )
						snake_case   :			Dict						    = int(items[1] )

						if type_id == 0:
												if "bias" in name:
																		if value.shape != feature_extractor.conv_layers[layer_id].conv.bias.data.shape:
																								raise ValueError(
																								    f'''{full_name} has size {value.shape}, but'''
																								    f''' {feature_extractor.conv_layers[layer_id].conv.bias.data.shape} was found.''' )
																		snake_case   :			Optional[int]						    = value
																		logger.info(f'''Feat extract conv layer {layer_id} was initialized from {full_name}.''' )
												elif "weight" in name:
																		if value.shape != feature_extractor.conv_layers[layer_id].conv.weight.data.shape:
																								raise ValueError(
																								    f'''{full_name} has size {value.shape}, but'''
																								    f''' {feature_extractor.conv_layers[layer_id].conv.weight.data.shape} was found.''' )
																		snake_case   :			Tuple						    = 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:
																		if value.shape != feature_extractor.conv_layers[layer_id].layer_norm.bias.data.shape:
																								raise ValueError(
																								    f'''{full_name} has size {value.shape}, but'''
																								    f''' {feature_extractor.conv_layers[layer_id].layer_norm.bias.data.shape} was found.''' )
																		snake_case   :			Dict						    = value
																		logger.info(f'''Feat extract layer norm weight of layer {layer_id} was initialized from {full_name}.''' )
												elif "weight" in name:
																		if value.shape != feature_extractor.conv_layers[layer_id].layer_norm.weight.data.shape:
																								raise ValueError(
																								    f'''{full_name} has size {value.shape}, but'''
																								    f''' {feature_extractor.conv_layers[layer_id].layer_norm.weight.data.shape} was found.''' )
																		snake_case   :			List[str]						    = value
																		logger.info(f'''Feat extract layer norm weight of layer {layer_id} was initialized from {full_name}.''' )
						else:
												unused_weights.append(lowerCamelCase_ )

@torch.no_grad()
def  __SCREAMING_SNAKE_CASE		(			lowerCamelCase_:							Optional[int]     ,		lowerCamelCase_:							str     ,		lowerCamelCase_:							List[str]=None     ,		lowerCamelCase_:							List[str]=None     ,		lowerCamelCase_:							Any=True ):
						"""simple docstring"""






						if config_path is not None:
												snake_case   :			Optional[int]						    = WavaVecaConformerConfig.from_pretrained(lowerCamelCase_     ,		hidden_act="swish" )
						else:
												snake_case   :			Union[str, Any]						    = WavaVecaConformerConfig()

						if "rope" in checkpoint_path:
												snake_case   :			Dict						    = "rotary"

						if is_finetuned:
												if dict_path:
																		snake_case   :			Optional[Any]						    = Dictionary.load(lowerCamelCase_ )

																		# important change bos & pad token id since CTC symbol is <pad> and
																		# not <s> as in fairseq
																		snake_case   :			List[Any]						    = target_dict.pad_index
																		snake_case   :			List[Any]						    = target_dict.bos_index
																		snake_case   :			str						    = target_dict.eos_index
																		snake_case   :			Optional[int]						    = len(target_dict.symbols )
																		snake_case   :			Dict						    = os.path.join(lowerCamelCase_     ,		"vocab.json" )
																		if not os.path.isdir(lowerCamelCase_ ):
																								logger.error("--pytorch_dump_folder_path ({}) should be a directory".format(lowerCamelCase_ ) )
																								return
																		os.makedirs(lowerCamelCase_     ,		exist_ok=lowerCamelCase_ )
																		snake_case   :			Union[str, Any]						    = target_dict.indices

																		# fairseq has the <pad> and <s> switched
																		snake_case   :			int						    = 0
																		snake_case   :			int						    = 1
																		with open(lowerCamelCase_     ,		"w"     ,		encoding="utf-8" ) as vocab_handle:
																								json.dump(lowerCamelCase_     ,		lowerCamelCase_ )
																		snake_case   :			List[str]						    = WavaVecaCTCTokenizer(
																		    lowerCamelCase_     ,		unk_token=target_dict.unk_word     ,		pad_token=target_dict.pad_word     ,		bos_token=target_dict.bos_word     ,		eos_token=target_dict.eos_word     ,		word_delimiter_token="|"     ,		do_lower_case=lowerCamelCase_     ,		)
																		snake_case   :			Any						    = True if config.feat_extract_norm == "layer" else False
																		snake_case   :			Optional[Any]						    = WavaVecaFeatureExtractor(
																		    feature_size=1     ,		sampling_rate=1_6_0_0_0     ,		padding_value=0     ,		do_normalize=lowerCamelCase_     ,		return_attention_mask=lowerCamelCase_     ,		)
																		snake_case   :			List[Any]						    = WavaVecaProcessor(feature_extractor=lowerCamelCase_     ,		tokenizer=lowerCamelCase_ )
																		processor.save_pretrained(lowerCamelCase_ )

												snake_case   :			Tuple						    = WavaVecaConformerForCTC(lowerCamelCase_ )
						else:
												snake_case   :			Union[str, Any]						    = WavaVecaConformerForPreTraining(lowerCamelCase_ )

						if is_finetuned:
												snake_case    ,			snake_case    ,			snake_case   :			str						    = fairseq.checkpoint_utils.load_model_ensemble_and_task(
												    [checkpoint_path]     ,		arg_overrides={"data": "/".join(dict_path.split("/" )[:-1] )} )
						else:
												snake_case   :			Union[str, Any]						    = argparse.Namespace(task="audio_pretraining" )
												snake_case   :			Optional[int]						    = fairseq.tasks.setup_task(lowerCamelCase_ )

												snake_case    ,			snake_case    ,			snake_case   :			Any						    = fairseq.checkpoint_utils.load_model_ensemble_and_task([checkpoint_path]     ,		task=lowerCamelCase_ )

						snake_case   :			List[Any]						    = model[0].eval()

						recursively_load_weights(lowerCamelCase_     ,		lowerCamelCase_     ,		not is_finetuned )

						hf_wavavec.save_pretrained(lowerCamelCase_ )


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_path', default=None, type=str, help='Path to hf config.json of model to convert')
			parser.add_argument(
			    '--not_finetuned', action='store_true', help='Whether the model to convert is a fine-tuned model or not'
			)
			A		=						parser.parse_args()
			convert_wavaveca_conformer_checkpoint(
			    args.checkpoint_path, args.pytorch_dump_folder_path, args.config_path, args.dict_path, not args.not_finetuned
			)








"""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       _a  ( SCREAMING_SNAKE_CASE__     ,	SCREAMING_SNAKE_CASE__     ,	unittest.TestCase):
			__magic_name__     					=     StableDiffusionPanoramaPipeline
			__magic_name__     					=     TEXT_TO_IMAGE_PARAMS
			__magic_name__     					=     TEXT_TO_IMAGE_BATCH_PARAMS
			__magic_name__     					=     TEXT_TO_IMAGE_IMAGE_PARAMS
			__magic_name__     					=     TEXT_TO_IMAGE_IMAGE_PARAMS



			def 				__lowercase					(			self							:				Optional[int]		)							->					int:
									torch.manual_seed(0		)
									snake_case   :			Optional[Any]						    = 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   ,	)
									snake_case   :			Tuple						    = DDIMScheduler()
									torch.manual_seed(0		)
									snake_case   :			List[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   ,	)
									torch.manual_seed(0		)
									snake_case   :			str						    = 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=1000   ,	)
									snake_case   :			int						    = CLIPTextModel(_lowercase		)
									snake_case   :			Union[str, Any]						    = CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"		)

									snake_case   :			Any						    = {
									    "unet": unet,
									    "scheduler": scheduler,
									    "vae": vae,
									    "text_encoder": text_encoder,
									    "tokenizer": tokenizer,
									    "safety_checker": None,
									    "feature_extractor": None,
									}
									return components



			def 				__lowercase					(			self							:				List[str]   ,	_lowercase							:				int   ,	_lowercase							:				Dict=0		)							->					Optional[int]:
									snake_case   :			Union[str, Any]						    = torch.manual_seed(_lowercase		)
									snake_case   :			str						    = {
									    "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 				__lowercase					(			self							:				str		)							->					List[str]:
									snake_case   :			Any						    = "cpu"  # ensure determinism for the device-dependent torch.Generator
									snake_case   :			Dict						    = self.get_dummy_components()
									snake_case   :			Any						    = StableDiffusionPanoramaPipeline(**_lowercase		)
									snake_case   :			List[Any]						    = sd_pipe.to(_lowercase		)
									sd_pipe.set_progress_bar_config(disable=_lowercase		)

									snake_case   :			Optional[Any]						    = self.get_dummy_inputs(_lowercase		)
									snake_case   :			Union[str, Any]						    = sd_pipe(**_lowercase		).images
									snake_case   :			List[Any]						    = image[0, -3:, -3:, -1]
									assert image.shape == (1, 64, 64, 3)

									snake_case   :			List[Any]						    = 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 				__lowercase					(			self							:				int		)							->					Union[str, Any]:
									super().test_inference_batch_consistent(batch_sizes=[1, 2]		)



			def 				__lowercase					(			self							:				Tuple		)							->					Tuple:
									super().test_inference_batch_single_identical(batch_size=2   ,	expected_max_diff=3.25E-3		)



			def 				__lowercase					(			self							:				Any		)							->					List[Any]:
									snake_case   :			Any						    = "cpu"  # ensure determinism for the device-dependent torch.Generator
									snake_case   :			Union[str, Any]						    = self.get_dummy_components()
									snake_case   :			Tuple						    = StableDiffusionPanoramaPipeline(**_lowercase		)
									snake_case   :			Tuple						    = sd_pipe.to(_lowercase		)
									sd_pipe.set_progress_bar_config(disable=_lowercase		)

									snake_case   :			List[str]						    = self.get_dummy_inputs(_lowercase		)
									snake_case   :			int						    = "french fries"
									snake_case   :			Union[str, Any]						    = sd_pipe(**_lowercase   ,	negative_prompt=_lowercase		)
									snake_case   :			str						    = output.images
									snake_case   :			Optional[Any]						    = image[0, -3:, -3:, -1]

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

									snake_case   :			Union[str, Any]						    = 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 				__lowercase					(			self							:				str		)							->					Any:
									snake_case   :			Optional[Any]						    = "cpu"  # ensure determinism for the device-dependent torch.Generator
									snake_case   :			List[Any]						    = self.get_dummy_components()
									snake_case   :			List[Any]						    = StableDiffusionPanoramaPipeline(**_lowercase		)
									snake_case   :			int						    = sd_pipe.to(_lowercase		)
									sd_pipe.set_progress_bar_config(disable=_lowercase		)

									snake_case   :			Tuple						    = self.get_dummy_inputs(_lowercase		)
									snake_case   :			str						    = sd_pipe(**_lowercase   ,	view_batch_size=2		)
									snake_case   :			Optional[Any]						    = output.images
									snake_case   :			int						    = image[0, -3:, -3:, -1]

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

									snake_case   :			str						    = 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 				__lowercase					(			self							:				int		)							->					Optional[Any]:
									snake_case   :			Dict						    = "cpu"  # ensure determinism for the device-dependent torch.Generator
									snake_case   :			List[str]						    = self.get_dummy_components()
									snake_case   :			Any						    = EulerAncestralDiscreteScheduler(
									    beta_start=0.00085   ,	beta_end=0.012   ,	beta_schedule="scaled_linear"		)
									snake_case   :			List[str]						    = StableDiffusionPanoramaPipeline(**_lowercase		)
									snake_case   :			List[str]						    = sd_pipe.to(_lowercase		)
									sd_pipe.set_progress_bar_config(disable=_lowercase		)

									snake_case   :			List[Any]						    = self.get_dummy_inputs(_lowercase		)
									snake_case   :			Optional[Any]						    = sd_pipe(**_lowercase		).images
									snake_case   :			int						    = image[0, -3:, -3:, -1]

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

									snake_case   :			Tuple						    = 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 				__lowercase					(			self							:				Tuple		)							->					Union[str, Any]:
									snake_case   :			str						    = "cpu"  # ensure determinism for the device-dependent torch.Generator
									snake_case   :			str						    = self.get_dummy_components()
									snake_case   :			Optional[int]						    = PNDMScheduler(
									    beta_start=0.00085   ,	beta_end=0.012   ,	beta_schedule="scaled_linear"   ,	skip_prk_steps=_lowercase		)
									snake_case   :			Dict						    = StableDiffusionPanoramaPipeline(**_lowercase		)
									snake_case   :			Dict						    = sd_pipe.to(_lowercase		)
									sd_pipe.set_progress_bar_config(disable=_lowercase		)

									snake_case   :			Dict						    = self.get_dummy_inputs(_lowercase		)
									snake_case   :			Optional[int]						    = sd_pipe(**_lowercase		).images
									snake_case   :			str						    = image[0, -3:, -3:, -1]

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

									snake_case   :			Optional[int]						    = 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       _a  ( unittest.TestCase):



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



			def 				__lowercase					(			self							:				Dict   ,	_lowercase							:				Any=0		)							->					Optional[Any]:
									snake_case   :			Any						    = torch.manual_seed(_lowercase		)
									snake_case   :			Any						    = {
									    "prompt": "a photo of the dolomites",
									    "generator": generator,
									    "num_inference_steps": 3,
									    "guidance_scale": 7.5,
									    "output_type": "numpy",
									}
									return inputs



			def 				__lowercase					(			self							:				Optional[Any]		)							->					Union[str, Any]:
									snake_case   :			List[Any]						    = "stabilityai/stable-diffusion-2-base"
									snake_case   :			Dict						    = DDIMScheduler.from_pretrained(_lowercase   ,	subfolder="scheduler"		)
									snake_case   :			Any						    = StableDiffusionPanoramaPipeline.from_pretrained(_lowercase   ,	scheduler=_lowercase   ,	safety_checker=_lowercase		)
									pipe.to(_lowercase		)
									pipe.set_progress_bar_config(disable=_lowercase		)
									pipe.enable_attention_slicing()

									snake_case   :			str						    = self.get_inputs()
									snake_case   :			List[Any]						    = pipe(**_lowercase		).images
									snake_case   :			Optional[int]						    = image[0, -3:, -3:, -1].flatten()

									assert image.shape == (1, 512, 2048, 3)

									snake_case   :			List[str]						    = np.array(
									    [
									        0.36968392,
									        0.27025372,
									        0.32446766,
									        0.28379387,
									        0.36363274,
									        0.30733347,
									        0.27100027,
									        0.27054125,
									        0.25536096,
									    ]		)

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



			def 				__lowercase					(			self							:				Tuple		)							->					List[str]:
									snake_case   :			List[str]						    = StableDiffusionPanoramaPipeline.from_pretrained(
									    "stabilityai/stable-diffusion-2-base"   ,	safety_checker=_lowercase		)
									snake_case   :			int						    = LMSDiscreteScheduler.from_config(pipe.scheduler.config		)
									pipe.to(_lowercase		)
									pipe.set_progress_bar_config(disable=_lowercase		)
									pipe.enable_attention_slicing()

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

									assert image.shape == (1, 512, 2048, 3)

									snake_case   :			List[str]						    = 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 				__lowercase					(			self							:				str		)							->					Any:
									snake_case   :			Any						    = 0

									def callback_fn(_lowercase							:				int   ,	_lowercase							:				int   ,	_lowercase							:				torch.FloatTensor		) -> None:
															snake_case   :			List[Any]						    = True
															nonlocal number_of_steps
															number_of_steps += 1
															if step == 1:
																					snake_case   :			List[Any]						    = latents.detach().cpu().numpy()
																					assert latents.shape == (1, 4, 64, 256)
																					snake_case   :			int						    = latents[0, -3:, -3:, -1]

																					snake_case   :			Tuple						    = np.array(
																					    [
																					        0.18681869,
																					        0.33907816,
																					        0.5361276,
																					        0.14432865,
																					        -0.02856611,
																					        -0.73941123,
																					        0.23397987,
																					        0.47322682,
																					        -0.37823164,
																					    ]		)
																					assert np.abs(latents_slice.flatten() - expected_slice		).max() < 5E-2
															elif step == 2:
																					snake_case   :			Union[str, Any]						    = latents.detach().cpu().numpy()
																					assert latents.shape == (1, 4, 64, 256)
																					snake_case   :			Tuple						    = latents[0, -3:, -3:, -1]

																					snake_case   :			Tuple						    = np.array(
																					    [
																					        0.18539645,
																					        0.33987248,
																					        0.5378559,
																					        0.14437142,
																					        -0.02455261,
																					        -0.7338317,
																					        0.23990755,
																					        0.47356272,
																					        -0.3786505,
																					    ]		)

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

									snake_case   :			Dict						    = False

									snake_case   :			List[Any]						    = "stabilityai/stable-diffusion-2-base"
									snake_case   :			List[str]						    = DDIMScheduler.from_pretrained(_lowercase   ,	subfolder="scheduler"		)
									snake_case   :			Optional[Any]						    = StableDiffusionPanoramaPipeline.from_pretrained(_lowercase   ,	scheduler=_lowercase   ,	safety_checker=_lowercase		)
									snake_case   :			Any						    = pipe.to(_lowercase		)
									pipe.set_progress_bar_config(disable=_lowercase		)
									pipe.enable_attention_slicing()

									snake_case   :			Tuple						    = self.get_inputs()
									pipe(**_lowercase   ,	callback=_lowercase   ,	callback_steps=1		)
									assert callback_fn.has_been_called
									assert number_of_steps == 3



			def 				__lowercase					(			self							:				Any		)							->					Tuple:
									torch.cuda.empty_cache()
									torch.cuda.reset_max_memory_allocated()
									torch.cuda.reset_peak_memory_stats()

									snake_case   :			Optional[int]						    = "stabilityai/stable-diffusion-2-base"
									snake_case   :			Tuple						    = DDIMScheduler.from_pretrained(_lowercase   ,	subfolder="scheduler"		)
									snake_case   :			Any						    = StableDiffusionPanoramaPipeline.from_pretrained(_lowercase   ,	scheduler=_lowercase   ,	safety_checker=_lowercase		)
									snake_case   :			Tuple						    = pipe.to(_lowercase		)
									pipe.set_progress_bar_config(disable=_lowercase		)
									pipe.enable_attention_slicing(1		)
									pipe.enable_sequential_cpu_offload()

									snake_case   :			Any						    = self.get_inputs()
									snake_case   :			Optional[int]						    = pipe(**_lowercase		)

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



