

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()
_SCREAMING_SNAKE_CASE  :						Dict         =    logging.get_logger(__name__)

_SCREAMING_SNAKE_CASE  :						Optional[Any]         =    {
    '''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''',
}
_SCREAMING_SNAKE_CASE  :						Optional[Any]         =    [
    '''lm_head''',
    '''quantizer.weight_proj''',
    '''quantizer.codevectors''',
    '''project_q''',
    '''project_hid''',
]
def        UpperCAmelCase_     (		_A   ,	_A   ,	_A   ,	_A   ,	_A       ):





     '''simple docstring'''




     for attribute in key.split('''.'''       ):
          SCREAMING_SNAKE_CASE__					=					getattr(_A   ,	_A       )

     if weight_type is not None:
          SCREAMING_SNAKE_CASE__					=					getattr(_A   ,	_A       ).shape
     else:
          SCREAMING_SNAKE_CASE__					=					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":
          SCREAMING_SNAKE_CASE__					=					value
     elif weight_type == "weight_g":
          SCREAMING_SNAKE_CASE__					=					value
     elif weight_type == "weight_v":
          SCREAMING_SNAKE_CASE__					=					value
     elif weight_type == "bias":
          SCREAMING_SNAKE_CASE__					=					value
     elif weight_type == "running_mean":
          SCREAMING_SNAKE_CASE__					=					value
     elif weight_type == "running_var":
          SCREAMING_SNAKE_CASE__					=					value
     elif weight_type == "num_batches_tracked":
          SCREAMING_SNAKE_CASE__					=					value
     elif weight_type == "inv_freq":
          SCREAMING_SNAKE_CASE__					=					value
     else:
          SCREAMING_SNAKE_CASE__					=					value

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





     '''simple docstring'''




     SCREAMING_SNAKE_CASE__					=					[]
     SCREAMING_SNAKE_CASE__					=					fairseq_model.state_dict()

     SCREAMING_SNAKE_CASE__					=					hf_model.wavaveca_conformer.feature_extractor

     for name, value in fairseq_dict.items():
          SCREAMING_SNAKE_CASE__					=					False
          if "conv_layers" in name:
               load_conv_layer(
                   _A   ,	_A   ,	_A   ,	_A   ,	hf_model.config.feat_extract_norm == '''group'''   ,	)
               SCREAMING_SNAKE_CASE__					=					True
          else:
               for key, mapped_key in MAPPING.items():
                    SCREAMING_SNAKE_CASE__					=					'''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]:
                         SCREAMING_SNAKE_CASE__					=					True
                         if "*" in mapped_key:
                              SCREAMING_SNAKE_CASE__					=					name.split(_A       )[0].split('''.'''       )[-2]
                              SCREAMING_SNAKE_CASE__					=					mapped_key.replace('''*'''   ,	_A       )
                         if "pos_bias_u" in name:
                              SCREAMING_SNAKE_CASE__					=					None
                         elif "pos_bias_v" in name:
                              SCREAMING_SNAKE_CASE__					=					None
                         elif "weight_g" in name:
                              SCREAMING_SNAKE_CASE__					=					'''weight_g'''
                         elif "weight_v" in name:
                              SCREAMING_SNAKE_CASE__					=					'''weight_v'''
                         elif "bias" in name:
                              SCREAMING_SNAKE_CASE__					=					'''bias'''
                         elif "weight" in name:
                              # TODO: don't match quantizer.weight_proj
                              SCREAMING_SNAKE_CASE__					=					'''weight'''
                         elif "running_mean" in name:
                              SCREAMING_SNAKE_CASE__					=					'''running_mean'''
                         elif "inv_freq" in name:
                              SCREAMING_SNAKE_CASE__					=					'''inv_freq'''
                         elif "running_var" in name:
                              SCREAMING_SNAKE_CASE__					=					'''running_var'''
                         elif "num_batches_tracked" in name:
                              SCREAMING_SNAKE_CASE__					=					'''num_batches_tracked'''
                         else:
                              SCREAMING_SNAKE_CASE__					=					None
                         set_recursively(_A   ,	_A   ,	_A   ,	_A   ,	_A       )
                    continue
          if not is_used:
               unused_weights.append(_A       )

     logger.warning(F'''Unused weights: {unused_weights}'''       )
def        UpperCAmelCase_     (		_A   ,	_A   ,	_A   ,	_A   ,	_A       ):





     '''simple docstring'''




     SCREAMING_SNAKE_CASE__					=					full_name.split('''conv_layers.'''       )[-1]
     SCREAMING_SNAKE_CASE__					=					name.split('''.'''       )
     SCREAMING_SNAKE_CASE__					=					int(items[0]       )
     SCREAMING_SNAKE_CASE__					=					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.'''       )
               SCREAMING_SNAKE_CASE__					=					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.'''       )
               SCREAMING_SNAKE_CASE__					=					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.'''       )
               SCREAMING_SNAKE_CASE__					=					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.'''       )
               SCREAMING_SNAKE_CASE__					=					value
               logger.info(F'''Feat extract layer norm weight of layer {layer_id} was initialized from {full_name}.'''       )
     else:
          unused_weights.append(_A       )
@torch.no_grad()
def        UpperCAmelCase_     (		_A   ,	_A   ,	_A=None   ,	_A=None   ,	_A=True       ):





     '''simple docstring'''




     if config_path is not None:
          SCREAMING_SNAKE_CASE__					=					WavaVecaConformerConfig.from_pretrained(_A   ,	hidden_act='''swish'''       )
     else:
          SCREAMING_SNAKE_CASE__					=					WavaVecaConformerConfig()

     if "rope" in checkpoint_path:
          SCREAMING_SNAKE_CASE__					=					'''rotary'''

     if is_finetuned:
          if dict_path:
               SCREAMING_SNAKE_CASE__					=					Dictionary.load(_A       )

               # important change bos & pad token id since CTC symbol is <pad> and
               # not <s> as in fairseq
               SCREAMING_SNAKE_CASE__					=					target_dict.pad_index
               SCREAMING_SNAKE_CASE__					=					target_dict.bos_index
               SCREAMING_SNAKE_CASE__					=					target_dict.eos_index
               SCREAMING_SNAKE_CASE__					=					len(target_dict.symbols       )
               SCREAMING_SNAKE_CASE__					=					os.path.join(_A   ,	'''vocab.json'''       )
               if not os.path.isdir(_A       ):
                    logger.error('''--pytorch_dump_folder_path ({}) should be a directory'''.format(_A       )       )
                    return
               os.makedirs(_A   ,	exist_ok=_A       )
               SCREAMING_SNAKE_CASE__					=					target_dict.indices

               # fairseq has the <pad> and <s> switched
               SCREAMING_SNAKE_CASE__					=					0
               SCREAMING_SNAKE_CASE__					=					1
               with open(_A   ,	'''w'''   ,	encoding='''utf-8'''       ) as vocab_handle:
                    json.dump(_A   ,	_A       )
               SCREAMING_SNAKE_CASE__					=					WavaVecaCTCTokenizer(
                   _A   ,	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=_A   ,	)
               SCREAMING_SNAKE_CASE__					=					True if config.feat_extract_norm == '''layer''' else False
               SCREAMING_SNAKE_CASE__					=					WavaVecaFeatureExtractor(
                   feature_size=1   ,	sampling_rate=1_60_00   ,	padding_value=0   ,	do_normalize=_A   ,	return_attention_mask=_A   ,	)
               SCREAMING_SNAKE_CASE__					=					WavaVecaProcessor(feature_extractor=_A   ,	tokenizer=_A       )
               processor.save_pretrained(_A       )

          SCREAMING_SNAKE_CASE__					=					WavaVecaConformerForCTC(_A       )
     else:
          SCREAMING_SNAKE_CASE__					=					WavaVecaConformerForPreTraining(_A       )

     if is_finetuned:
          SCREAMING_SNAKE_CASE__,SCREAMING_SNAKE_CASE__,SCREAMING_SNAKE_CASE__					=					fairseq.checkpoint_utils.load_model_ensemble_and_task(
              [checkpoint_path]   ,	arg_overrides={'''data''': '''/'''.join(dict_path.split('''/'''       )[:-1]       )}       )
     else:
          SCREAMING_SNAKE_CASE__					=					argparse.Namespace(task='''audio_pretraining'''       )
          SCREAMING_SNAKE_CASE__					=					fairseq.tasks.setup_task(_A       )

          SCREAMING_SNAKE_CASE__,SCREAMING_SNAKE_CASE__,SCREAMING_SNAKE_CASE__					=					fairseq.checkpoint_utils.load_model_ensemble_and_task([checkpoint_path]   ,	task=_A       )

     SCREAMING_SNAKE_CASE__					=					model[0].eval()

     recursively_load_weights(_A   ,	_A   ,	not is_finetuned       )

     hf_wavavec.save_pretrained(_A       )


if __name__ == "__main__":
       _SCREAMING_SNAKE_CASE  :						List[Any]         =    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'''
       )
       _SCREAMING_SNAKE_CASE  :						Union[str, Any]         =    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
       )





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__	( A__		,  A__		,  unittest.TestCase ):






  """simple docstring"""


  a     =      StableDiffusionPanoramaPipeline
  a     =      TEXT_TO_IMAGE_PARAMS
  a     =      TEXT_TO_IMAGE_BATCH_PARAMS
  a     =      TEXT_TO_IMAGE_IMAGE_PARAMS
  a     =      TEXT_TO_IMAGE_IMAGE_PARAMS






  def    lowercase_	(							self		:	Union[str, Any]      )	->							List[str]:
       torch.manual_seed(0      )
       SCREAMING_SNAKE_CASE__					=					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			,     )
       SCREAMING_SNAKE_CASE__					=					DDIMScheduler()
       torch.manual_seed(0      )
       SCREAMING_SNAKE_CASE__					=					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      )
       SCREAMING_SNAKE_CASE__					=					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			,     )
       SCREAMING_SNAKE_CASE__					=					CLIPTextModel(__lowerCamelCase      )
       SCREAMING_SNAKE_CASE__					=					CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''      )

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






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

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

       SCREAMING_SNAKE_CASE__					=					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		:	Any      )	->							Optional[Any]:
       super().test_inference_batch_consistent(batch_sizes=[1, 2]      )






  def    lowercase_	(							self		:	List[Any]      )	->							List[Any]:
       super().test_inference_batch_single_identical(batch_size=2			,     expected_max_diff=3.25e-3      )






  def    lowercase_	(							self		:	Union[str, Any]      )	->							str:
       SCREAMING_SNAKE_CASE__					=					'''cpu'''  # ensure determinism for the device-dependent torch.Generator
       SCREAMING_SNAKE_CASE__					=					self.get_dummy_components()
       SCREAMING_SNAKE_CASE__					=					StableDiffusionPanoramaPipeline(**__lowerCamelCase      )
       SCREAMING_SNAKE_CASE__					=					sd_pipe.to(__lowerCamelCase      )
       sd_pipe.set_progress_bar_config(disable=__lowerCamelCase      )

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

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

       SCREAMING_SNAKE_CASE__					=					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		:	Tuple      )	->							Dict:
       SCREAMING_SNAKE_CASE__					=					'''cpu'''  # ensure determinism for the device-dependent torch.Generator
       SCREAMING_SNAKE_CASE__					=					self.get_dummy_components()
       SCREAMING_SNAKE_CASE__					=					StableDiffusionPanoramaPipeline(**__lowerCamelCase      )
       SCREAMING_SNAKE_CASE__					=					sd_pipe.to(__lowerCamelCase      )
       sd_pipe.set_progress_bar_config(disable=__lowerCamelCase      )

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

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

       SCREAMING_SNAKE_CASE__					=					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		:	Optional[Any]      )	->							Any:
       SCREAMING_SNAKE_CASE__					=					'''cpu'''  # ensure determinism for the device-dependent torch.Generator
       SCREAMING_SNAKE_CASE__					=					self.get_dummy_components()
       SCREAMING_SNAKE_CASE__					=					EulerAncestralDiscreteScheduler(
           beta_start=0.00085			,     beta_end=0.012			,     beta_schedule='''scaled_linear'''      )
       SCREAMING_SNAKE_CASE__					=					StableDiffusionPanoramaPipeline(**__lowerCamelCase      )
       SCREAMING_SNAKE_CASE__					=					sd_pipe.to(__lowerCamelCase      )
       sd_pipe.set_progress_bar_config(disable=__lowerCamelCase      )

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

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

       SCREAMING_SNAKE_CASE__					=					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		:	Optional[int]      )	->							List[str]:
       SCREAMING_SNAKE_CASE__					=					'''cpu'''  # ensure determinism for the device-dependent torch.Generator
       SCREAMING_SNAKE_CASE__					=					self.get_dummy_components()
       SCREAMING_SNAKE_CASE__					=					PNDMScheduler(
           beta_start=0.00085			,     beta_end=0.012			,     beta_schedule='''scaled_linear'''			,     skip_prk_steps=__lowerCamelCase      )
       SCREAMING_SNAKE_CASE__					=					StableDiffusionPanoramaPipeline(**__lowerCamelCase      )
       SCREAMING_SNAKE_CASE__					=					sd_pipe.to(__lowerCamelCase      )
       sd_pipe.set_progress_bar_config(disable=__lowerCamelCase      )

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

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

       SCREAMING_SNAKE_CASE__					=					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 ):






  """simple docstring"""


  def    lowercase_	(							self		:	Tuple      )	->							str:
       super().tearDown()
       gc.collect()
       torch.cuda.empty_cache()






  def    lowercase_	(							self		:	Any			,     __lowerCamelCase		:	Dict=0      )	->							Dict:
       SCREAMING_SNAKE_CASE__					=					torch.manual_seed(__lowerCamelCase      )
       SCREAMING_SNAKE_CASE__					=					{
           '''prompt''': '''a photo of the dolomites''',
           '''generator''': generator,
           '''num_inference_steps''': 3,
           '''guidance_scale''': 7.5,
           '''output_type''': '''numpy''',
       }
       return inputs






  def    lowercase_	(							self		:	Union[str, Any]      )	->							List[Any]:
       SCREAMING_SNAKE_CASE__					=					'''stabilityai/stable-diffusion-2-base'''
       SCREAMING_SNAKE_CASE__					=					DDIMScheduler.from_pretrained(__lowerCamelCase			,     subfolder='''scheduler'''      )
       SCREAMING_SNAKE_CASE__					=					StableDiffusionPanoramaPipeline.from_pretrained(__lowerCamelCase			,     scheduler=__lowerCamelCase			,     safety_checker=__lowerCamelCase      )
       pipe.to(__lowerCamelCase      )
       pipe.set_progress_bar_config(disable=__lowerCamelCase      )
       pipe.enable_attention_slicing()

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

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

       SCREAMING_SNAKE_CASE__					=					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		:	Optional[Any]      )	->							Dict:
       SCREAMING_SNAKE_CASE__					=					StableDiffusionPanoramaPipeline.from_pretrained(
           '''stabilityai/stable-diffusion-2-base'''			,     safety_checker=__lowerCamelCase      )
       SCREAMING_SNAKE_CASE__					=					LMSDiscreteScheduler.from_config(pipe.scheduler.config      )
       pipe.to(__lowerCamelCase      )
       pipe.set_progress_bar_config(disable=__lowerCamelCase      )
       pipe.enable_attention_slicing()

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

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

       SCREAMING_SNAKE_CASE__					=					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		:	List[str]      )	->							List[str]:
       SCREAMING_SNAKE_CASE__					=					0

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

                 SCREAMING_SNAKE_CASE__					=					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:
                 SCREAMING_SNAKE_CASE__					=					latents.detach().cpu().numpy()
                 assert latents.shape == (1, 4, 64, 256)
                 SCREAMING_SNAKE_CASE__					=					latents[0, -3:, -3:, -1]

                 SCREAMING_SNAKE_CASE__					=					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

       SCREAMING_SNAKE_CASE__					=					False

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

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







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

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

       SCREAMING_SNAKE_CASE__					=					self.get_inputs()
       SCREAMING_SNAKE_CASE__					=					pipe(**__lowerCamelCase      )

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


