



"""simple docstring"""







import argparse
import collections
import json
from pathlib import Path

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

from transformers import (
    MobileViTImageProcessor,
    MobileViTVaConfig,
    MobileViTVaForImageClassification,
    MobileViTVaForSemanticSegmentation,
)
from transformers.utils import logging


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

def  __SCREAMING_SNAKE_CASE		(			lowerCamelCase_:							int ):
						"""simple docstring"""






						print("Loading config file..." )

						def flatten_yaml_as_dict(lowerCamelCase_:							List[str]     ,		lowerCamelCase_:							List[Any]=""     ,		lowerCamelCase_:							Optional[int]="." ):
												snake_case   :			Optional[int]						    = []
												for k, v in d.items():
																		snake_case   :			str						    = parent_key + sep + k if parent_key else k
																		if isinstance(lowerCamelCase_     ,		collections.abc.MutableMapping ):
																								items.extend(flatten_yaml_as_dict(lowerCamelCase_     ,		lowerCamelCase_     ,		sep=lowerCamelCase_ ).items() )
																		else:
																								items.append((new_key, v) )
												return dict(lowerCamelCase_ )

						snake_case   :			List[str]						    = argparse.Namespace()
						with open(lowerCamelCase_     ,		"r" ) as yaml_file:
												try:
																		snake_case   :			Tuple						    = yaml.load(lowerCamelCase_     ,		Loader=yaml.FullLoader )

																		snake_case   :			Dict						    = flatten_yaml_as_dict(lowerCamelCase_ )
																		for k, v in flat_cfg.items():
																								setattr(lowerCamelCase_     ,		lowerCamelCase_     ,		lowerCamelCase_ )
												except yaml.YAMLError as exc:
																		logger.error("Error while loading config file: {}. Error message: {}".format(lowerCamelCase_     ,		str(lowerCamelCase_ ) ) )
						return config

def  __SCREAMING_SNAKE_CASE		(			lowerCamelCase_:							Dict     ,		lowerCamelCase_:							List[Any] ):
						"""simple docstring"""






						snake_case   :			List[str]						    = MobileViTVaConfig()

						snake_case   :			str						    = False

						# dataset
						if task_name.startswith("imagenet1k_" ):
												snake_case   :			int						    = 1_0_0_0
												if int(task_name.strip().split("_" )[-1] ) == 3_8_4:
																		snake_case   :			int						    = 3_8_4
												else:
																		snake_case   :			Union[str, Any]						    = 2_5_6
												snake_case   :			str						    = "imagenet-1k-id2label.json"
						elif task_name.startswith("imagenet21k_to_1k_" ):
												snake_case   :			Optional[int]						    = 2_1_0_0_0
												if int(task_name.strip().split("_" )[-1] ) == 3_8_4:
																		snake_case   :			Dict						    = 3_8_4
												else:
																		snake_case   :			str						    = 2_5_6
												snake_case   :			Optional[Any]						    = "imagenet-22k-id2label.json"
						elif task_name.startswith("ade20k_" ):
												snake_case   :			int						    = 1_5_1
												snake_case   :			Optional[int]						    = 5_1_2
												snake_case   :			Union[str, Any]						    = "ade20k-id2label.json"
												snake_case   :			Optional[int]						    = True
						elif task_name.startswith("voc_" ):
												snake_case   :			Union[str, Any]						    = 2_1
												snake_case   :			Tuple						    = 5_1_2
												snake_case   :			Optional[int]						    = "pascal-voc-id2label.json"
												snake_case   :			Any						    = True

						# orig_config
						snake_case   :			Optional[Any]						    = load_orig_config_file(lowerCamelCase_ )
						assert getattr(lowerCamelCase_     ,		"model.classification.name"     ,		-1 ) == "mobilevit_v2", "Invalid model"
						snake_case   :			Union[str, Any]						    = getattr(lowerCamelCase_     ,		"model.classification.mitv2.width_multiplier"     ,		1.0 )
						assert (
						    getattr(lowerCamelCase_     ,		"model.classification.mitv2.attn_norm_layer"     ,		-1 ) == "layer_norm_2d"
						), "Norm layers other than layer_norm_2d is not supported"
						snake_case   :			Optional[int]						    = getattr(lowerCamelCase_     ,		"model.classification.activation.name"     ,		"swish" )
						# config.image_size == getattr(orig_config,  'sampler.bs.crop_size_width', 256)

						if is_segmentation_model:
												snake_case   :			List[str]						    = getattr(lowerCamelCase_     ,		"model.segmentation.output_stride"     ,		1_6 )
												if "_deeplabv3" in task_name:
																		snake_case   :			List[Any]						    = getattr(lowerCamelCase_     ,		"model.segmentation.deeplabv3.aspp_rates"     ,		[1_2, 2_4, 3_6] )
																		snake_case   :			Optional[int]						    = getattr(lowerCamelCase_     ,		"model.segmentation.deeplabv3.aspp_out_channels"     ,		5_1_2 )
																		snake_case   :			str						    = getattr(lowerCamelCase_     ,		"model.segmentation.deeplabv3.aspp_dropout"     ,		0.1 )

    # id2label
						snake_case   :			str						    = "huggingface/label-files"
						snake_case   :			List[Any]						    = json.load(open(hf_hub_download(lowerCamelCase_     ,		lowerCamelCase_     ,		repo_type="dataset" )     ,		"r" ) )
						snake_case   :			Any						    = {int(lowerCamelCase_ ): v for k, v in idalabel.items()}
						snake_case   :			int						    = idalabel
						snake_case   :			Any						    = {v: k for k, v in idalabel.items()}

						return config

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






						snake_case   :			Optional[int]						    = dct.pop(lowerCamelCase_ )
						snake_case   :			List[str]						    = val

def  __SCREAMING_SNAKE_CASE		(			lowerCamelCase_:							List[Any]     ,		lowerCamelCase_:							List[Any]=False ):
						"""simple docstring"""






						if base_model:
												snake_case   :			Dict						    = ""
						else:
												snake_case   :			Any						    = "mobilevitv2."

						snake_case   :			List[Any]						    = []
						for k in state_dict.keys():
												if k[:8] == "encoder.":
																		snake_case   :			Any						    = k[8:]
												else:
																		snake_case   :			Tuple						    = k

												if ".block." in k:
																		snake_case   :			Any						    = k_new.replace(".block."     ,		"." )
												if ".conv." in k:
																		snake_case   :			Any						    = k_new.replace(".conv."     ,		".convolution." )
												if ".norm." in k:
																		snake_case   :			Any						    = k_new.replace(".norm."     ,		".normalization." )

												if "conv_1." in k:
																		snake_case   :			Dict						    = k_new.replace("conv_1."     ,		f'''{model_prefix}conv_stem.''' )
												for i in [1, 2]:
																		if f'''layer_{i}.''' in k:
																								snake_case   :			int						    = k_new.replace(f'''layer_{i}.'''     ,		f'''{model_prefix}encoder.layer.{i-1}.layer.''' )
												if ".exp_1x1." in k:
																		snake_case   :			Any						    = k_new.replace(".exp_1x1."     ,		".expand_1x1." )
												if ".red_1x1." in k:
																		snake_case   :			Optional[int]						    = k_new.replace(".red_1x1."     ,		".reduce_1x1." )

												for i in [3, 4, 5]:
																		if f'''layer_{i}.0.''' in k:
																								snake_case   :			Optional[int]						    = k_new.replace(f'''layer_{i}.0.'''     ,		f'''{model_prefix}encoder.layer.{i-1}.downsampling_layer.''' )
																		if f'''layer_{i}.1.local_rep.0.''' in k:
																								snake_case   :			List[str]						    = k_new.replace(f'''layer_{i}.1.local_rep.0.'''     ,		f'''{model_prefix}encoder.layer.{i-1}.conv_kxk.''' )
																		if f'''layer_{i}.1.local_rep.1.''' in k:
																								snake_case   :			str						    = k_new.replace(f'''layer_{i}.1.local_rep.1.'''     ,		f'''{model_prefix}encoder.layer.{i-1}.conv_1x1.''' )

												for i in [3, 4, 5]:
																		if i == 3:
																								snake_case   :			Union[str, Any]						    = [0, 1]
																		elif i == 4:
																								snake_case   :			List[Any]						    = [0, 1, 2, 3]
																		elif i == 5:
																								snake_case   :			List[str]						    = [0, 1, 2]

																		for j in j_in:
																								if f'''layer_{i}.1.global_rep.{j}.''' in k:
																														snake_case   :			Dict						    = k_new.replace(
																														    f'''layer_{i}.1.global_rep.{j}.'''     ,		f'''{model_prefix}encoder.layer.{i-1}.transformer.layer.{j}.''' )
																		if f'''layer_{i}.1.global_rep.{j+1}.''' in k:
																								snake_case   :			str						    = k_new.replace(
																								    f'''layer_{i}.1.global_rep.{j+1}.'''     ,		f'''{model_prefix}encoder.layer.{i-1}.layernorm.''' )

																		if f'''layer_{i}.1.conv_proj.''' in k:
																								snake_case   :			int						    = k_new.replace(f'''layer_{i}.1.conv_proj.'''     ,		f'''{model_prefix}encoder.layer.{i-1}.conv_projection.''' )

												if "pre_norm_attn.0." in k:
																		snake_case   :			List[Any]						    = k_new.replace("pre_norm_attn.0."     ,		"layernorm_before." )
												if "pre_norm_attn.1." in k:
																		snake_case   :			int						    = k_new.replace("pre_norm_attn.1."     ,		"attention." )
												if "pre_norm_ffn.0." in k:
																		snake_case   :			Union[str, Any]						    = k_new.replace("pre_norm_ffn.0."     ,		"layernorm_after." )
												if "pre_norm_ffn.1." in k:
																		snake_case   :			Union[str, Any]						    = k_new.replace("pre_norm_ffn.1."     ,		"ffn.conv1." )
												if "pre_norm_ffn.3." in k:
																		snake_case   :			Optional[int]						    = k_new.replace("pre_norm_ffn.3."     ,		"ffn.conv2." )

												if "classifier.1." in k:
																		snake_case   :			Optional[Any]						    = k_new.replace("classifier.1."     ,		"classifier." )

												if "seg_head." in k:
																		snake_case   :			Tuple						    = k_new.replace("seg_head."     ,		"segmentation_head." )
												if ".aspp_layer." in k:
																		snake_case   :			Optional[int]						    = k_new.replace(".aspp_layer."     ,		"." )
												if ".aspp_pool." in k:
																		snake_case   :			Any						    = k_new.replace(".aspp_pool."     ,		"." )

												rename_keys.append((k, k_new) )
						return rename_keys

def  __SCREAMING_SNAKE_CASE		(			lowerCamelCase_:							List[Any] ):
						"""simple docstring"""






						snake_case   :			str						    = []
						for k in state_dict.keys():
												if k.startswith("seg_head.aux_head." ):
																		keys_to_ignore.append(lowerCamelCase_ )
						for k in keys_to_ignore:
												state_dict.pop(lowerCamelCase_     ,		lowerCamelCase_ )

def  __SCREAMING_SNAKE_CASE		(			):
						"""simple docstring"""






						snake_case   :			Union[str, Any]						    = "http://images.cocodataset.org/val2017/000000039769.jpg"
						# url = "https://cdn.britannica.com/86/141086-050-9D7C75EE/Gulfstream-G450-business-jet-passengers.jpg"
						snake_case   :			Any						    = Image.open(requests.get(lowerCamelCase_     ,		stream=lowerCamelCase_ ).raw )
						return im

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






						snake_case   :			Union[str, Any]						    = get_mobilevitva_config(lowerCamelCase_     ,		lowerCamelCase_ )

						# load original state_dict
						snake_case   :			Tuple						    = torch.load(lowerCamelCase_     ,		map_location="cpu" )

						# load huggingface model
						if task_name.startswith("ade20k_" ) or task_name.startswith("voc_" ):
												snake_case   :			Union[str, Any]						    = MobileViTVaForSemanticSegmentation(lowerCamelCase_ ).eval()
												snake_case   :			Any						    = False
						else:
												snake_case   :			Any						    = MobileViTVaForImageClassification(lowerCamelCase_ ).eval()
												snake_case   :			str						    = False

						# remove and rename some keys of load the original model
						snake_case   :			List[str]						    = checkpoint
						remove_unused_keys(lowerCamelCase_ )
						snake_case   :			Optional[int]						    = create_rename_keys(lowerCamelCase_     ,		base_model=lowerCamelCase_ )
						for rename_key_src, rename_key_dest in rename_keys:
												rename_key(lowerCamelCase_     ,		lowerCamelCase_     ,		lowerCamelCase_ )

						# load modified state_dict
						model.load_state_dict(lowerCamelCase_ )

						# Check outputs on an image, prepared by MobileViTImageProcessor
						snake_case   :			str						    = MobileViTImageProcessor(crop_size=config.image_size     ,		size=config.image_size + 3_2 )
						snake_case   :			Optional[int]						    = image_processor(images=prepare_img()     ,		return_tensors="pt" )
						snake_case   :			Dict						    = model(**lowerCamelCase_ )

						# verify classification model
						if task_name.startswith("imagenet" ):
												snake_case   :			str						    = outputs.logits
												snake_case   :			Optional[int]						    = logits.argmax(-1 ).item()
												print("Predicted class:"     ,		model.config.idalabel[predicted_class_idx] )
												if task_name.startswith("imagenet1k_256" ) and config.width_multiplier == 1.0:
																		# expected_logits for base variant
																		snake_case   :			Dict						    = torch.tensor([-1.6336e00, -7.3204e-02, -5.1883e-01] )
																		assert torch.allclose(logits[0, :3]     ,		lowerCamelCase_     ,		atol=1e-4 )

						Path(lowerCamelCase_ ).mkdir(exist_ok=lowerCamelCase_ )
						print(f'''Saving model {task_name} to {pytorch_dump_folder_path}''' )
						model.save_pretrained(lowerCamelCase_ )
						print(f'''Saving image processor to {pytorch_dump_folder_path}''' )
						image_processor.save_pretrained(lowerCamelCase_ )


if __name__ == "__main__":
			A		=						argparse.ArgumentParser()
			# Required parameters
			parser.add_argument(
			    '--task',
			    default='imagenet1k_256',
			    type=str,
			    help=(
			        'Name of the task for which the MobileViTV2 model you\'d like to convert is trained on . '
			        '\n                Classification (ImageNet-1k)\n                    - MobileViTV2 (256x256) : imagenet1k_256\n                    - MobileViTV2 (Trained on 256x256 and Finetuned on 384x384) : imagenet1k_384\n                    - MobileViTV2 (Trained on ImageNet-21k and Finetuned on ImageNet-1k 256x256) :\n                      imagenet21k_to_1k_256\n                    - MobileViTV2 (Trained on ImageNet-21k, Finetuned on ImageNet-1k 256x256, and Finetuned on\n                      ImageNet-1k 384x384) : imagenet21k_to_1k_384\n                Segmentation\n                    - ADE20K Dataset : ade20k_deeplabv3\n                    - Pascal VOC 2012 Dataset: voc_deeplabv3\n            '
			    ),
			    choices=[
			        'imagenet1k_256',
			        'imagenet1k_384',
			        'imagenet21k_to_1k_256',
			        'imagenet21k_to_1k_384',
			        'ade20k_deeplabv3',
			        'voc_deeplabv3',
			    ],
			)

			parser.add_argument(
			    '--orig_checkpoint_path', required=True, type=str, help='Path to the original state dict (.pt file).'
			)
			parser.add_argument('--orig_config_path', required=True, type=str, help='Path to the original config file.')
			parser.add_argument(
			    '--pytorch_dump_folder_path', required=True, type=str, help='Path to the output PyTorch model directory.'
			)

			A		=						parser.parse_args()
			convert_mobilevitva_checkpoint(
			    args.task, args.orig_checkpoint_path, args.orig_config_path, args.pytorch_dump_folder_path
			)








"""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



