


'''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   __lowerCAmelCase						(lowercase_  ,     lowercase_  ,     unittest.TestCase     ):





      '''simple docstring'''






      lowerCAmelCase__					:     Optional[int]  					=							StableDiffusionPanoramaPipeline
      lowerCAmelCase__					:     Optional[int]  					=							TEXT_TO_IMAGE_PARAMS
      lowerCAmelCase__					:     Union[str, Any]  					=							TEXT_TO_IMAGE_BATCH_PARAMS
      lowerCAmelCase__					:     Optional[Any]  					=							TEXT_TO_IMAGE_IMAGE_PARAMS
      lowerCAmelCase__					:     List[str]  					=							TEXT_TO_IMAGE_IMAGE_PARAMS



      def       UpperCamelCase__		(self    :    Dict ):

             '''simple docstring'''







             torch.manual_seed(0 )
             lowercase__      			= 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   ,						)
             lowercase__      			= DDIMScheduler()
             torch.manual_seed(0 )
             lowercase__      			= 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 )
             lowercase__      			= 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   ,						)
             lowercase__      			= CLIPTextModel(UpperCamelCase )
             lowercase__      			= CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip''' )

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



      def       UpperCamelCase__		(self    :    int   ,						UpperCamelCase    :    int   ,						UpperCamelCase    :    List[str]=0 ):

             '''simple docstring'''







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



      def       UpperCamelCase__		(self    :    Any ):

             '''simple docstring'''







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

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

             lowercase__      			= np.array([0.61_86, 0.53_74, 0.49_15, 0.41_35, 0.41_14, 0.45_63, 0.51_28, 0.49_77, 0.47_57] )

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



      def       UpperCamelCase__		(self    :    List[Any] ):

             '''simple docstring'''







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



      def       UpperCamelCase__		(self    :    List[str] ):

             '''simple docstring'''







             super().test_inference_batch_single_identical(batch_size=2   ,						expected_max_diff=3.25E-3 )



      def       UpperCamelCase__		(self    :    List[Any] ):

             '''simple docstring'''







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

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

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

             lowercase__      			= np.array([0.61_87, 0.53_75, 0.49_15, 0.41_36, 0.41_14, 0.45_63, 0.51_28, 0.49_76, 0.47_57] )

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



      def       UpperCamelCase__		(self    :    Optional[Any] ):

             '''simple docstring'''







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

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

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

             lowercase__      			= np.array([0.61_87, 0.53_75, 0.49_15, 0.41_36, 0.41_14, 0.45_63, 0.51_28, 0.49_76, 0.47_57] )

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



      def       UpperCamelCase__		(self    :    List[Any] ):

             '''simple docstring'''







             lowercase__      			= '''cpu'''  # ensure determinism for the device-dependent torch.Generator
             lowercase__      			= self.get_dummy_components()
             lowercase__      			= EulerAncestralDiscreteScheduler(
                 beta_start=0.0_00_85   ,						beta_end=0.0_12   ,						beta_schedule='''scaled_linear''' )
             lowercase__      			= StableDiffusionPanoramaPipeline(**UpperCamelCase )
             lowercase__      			= sd_pipe.to(UpperCamelCase )
             sd_pipe.set_progress_bar_config(disable=UpperCamelCase )

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

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

             lowercase__      			= np.array([0.40_24, 0.65_10, 0.49_01, 0.53_78, 0.58_13, 0.56_22, 0.47_95, 0.44_67, 0.49_52] )

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






      def       UpperCamelCase__		(self    :    List[str] ):

             '''simple docstring'''







             lowercase__      			= '''cpu'''  # ensure determinism for the device-dependent torch.Generator
             lowercase__      			= self.get_dummy_components()
             lowercase__      			= PNDMScheduler(
                 beta_start=0.0_00_85   ,						beta_end=0.0_12   ,						beta_schedule='''scaled_linear'''   ,						skip_prk_steps=UpperCamelCase )
             lowercase__      			= StableDiffusionPanoramaPipeline(**UpperCamelCase )
             lowercase__      			= sd_pipe.to(UpperCamelCase )
             sd_pipe.set_progress_bar_config(disable=UpperCamelCase )

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

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

             lowercase__      			= np.array([0.63_91, 0.62_91, 0.48_61, 0.51_34, 0.55_52, 0.45_78, 0.50_32, 0.50_23, 0.45_39] )

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






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





      '''simple docstring'''






      def       UpperCamelCase__		(self    :    List[Any] ):

             '''simple docstring'''







             super().tearDown()
             gc.collect()
             torch.cuda.empty_cache()



      def       UpperCamelCase__		(self    :    Dict   ,						UpperCamelCase    :    str=0 ):

             '''simple docstring'''







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



      def       UpperCamelCase__		(self    :    str ):

             '''simple docstring'''







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

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

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

             lowercase__      			= np.array(
                 [
                     0.36_96_83_92,
                     0.27_02_53_72,
                     0.32_44_67_66,
                     0.28_37_93_87,
                     0.36_36_32_74,
                     0.30_73_33_47,
                     0.27_10_00_27,
                     0.27_05_41_25,
                     0.25_53_60_96,
                 ] )

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



      def       UpperCamelCase__		(self    :    List[str] ):

             '''simple docstring'''







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

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

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

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

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



      def       UpperCamelCase__		(self    :    List[str] ):

             '''simple docstring'''







             lowercase__      			= 0

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

                           lowercase__      			= np.array(
                               [
                                   0.18_68_18_69,
                                   0.33_90_78_16,
                                   0.5_36_12_76,
                                   0.14_43_28_65,
                                   -0.02_85_66_11,
                                   -0.73_94_11_23,
                                   0.23_39_79_87,
                                   0.47_32_26_82,
                                   -0.37_82_31_64,
                               ] )
                           assert np.abs(latents_slice.flatten() - expected_slice ).max() < 5E-2
                    elif step == 2:
                           lowercase__      			= latents.detach().cpu().numpy()
                           assert latents.shape == (1, 4, 64, 256)
                           lowercase__      			= latents[0, -3:, -3:, -1]

                           lowercase__      			= np.array(
                               [
                                   0.18_53_96_45,
                                   0.33_98_72_48,
                                   0.5_37_85_59,
                                   0.14_43_71_42,
                                   -0.02_45_52_61,
                                   -0.7_33_83_17,
                                   0.23_99_07_55,
                                   0.47_35_62_72,
                                   -0.3_78_65_05,
                               ] )

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

             lowercase__      			= False

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

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






      def       UpperCamelCase__		(self    :    Optional[int] ):

             '''simple docstring'''







             torch.cuda.empty_cache()
             torch.cuda.reset_max_memory_allocated()
             torch.cuda.reset_peak_memory_stats()

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

             lowercase__      			= self.get_inputs()
             lowercase__      			= pipe(**UpperCamelCase )

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








'''simple docstring'''

import argparse

import torch

from transformers import BertForMaskedLM


if __name__ == "__main__":
 lowerCamelCase							:   List[str]  						=	argparse.ArgumentParser(
     description=(
         'Extraction some layers of the full BertForMaskedLM or RObertaForMaskedLM for Transfer Learned'
         ' Distillation'
     )
 )
 parser.add_argument('--model_type', default='bert', choices=['bert'])
 parser.add_argument('--model_name', default='bert-base-uncased', type=str)
 parser.add_argument('--dump_checkpoint', default='serialization_dir/tf_bert-base-uncased_0247911.pth', type=str)
 parser.add_argument('--vocab_transform', action='store_true')
 lowerCamelCase							:   str  						=	parser.parse_args()

 if args.model_type == "bert":
  lowerCamelCase							:   List[Any]  						=	BertForMaskedLM.from_pretrained(args.model_name)
  lowerCamelCase							:   Any  						=	'bert'
 else:
  raise ValueError('args.model_type should be "bert".')

 lowerCamelCase							:   int  						=	model.state_dict()
 lowerCamelCase							:   int  						=	{}

 for w in ["word_embeddings", "position_embeddings"]:
  lowerCamelCase							:   Any  						=	state_dict[f"""{prefix}.embeddings.{w}.weight"""]
 for w in ["weight", "bias"]:
  lowerCamelCase							:   Any  						=	state_dict[f"""{prefix}.embeddings.LayerNorm.{w}"""]

 lowerCamelCase							:   Tuple  						=	0
 for teacher_idx in [0, 2, 4, 7, 9, 11]:
  for w in ["weight", "bias"]:
   lowerCamelCase							:   Dict  						=	state_dict[
       f"""{prefix}.encoder.layer.{teacher_idx}.attention.self.query.{w}"""
   ]
   lowerCamelCase							:   Dict  						=	state_dict[
       f"""{prefix}.encoder.layer.{teacher_idx}.attention.self.key.{w}"""
   ]
   lowerCamelCase							:   List[Any]  						=	state_dict[
       f"""{prefix}.encoder.layer.{teacher_idx}.attention.self.value.{w}"""
   ]

   lowerCamelCase							:   Tuple  						=	state_dict[
       f"""{prefix}.encoder.layer.{teacher_idx}.attention.output.dense.{w}"""
   ]
   lowerCamelCase							:   Optional[int]  						=	state_dict[
       f"""{prefix}.encoder.layer.{teacher_idx}.attention.output.LayerNorm.{w}"""
   ]

   lowerCamelCase							:   Optional[Any]  						=	state_dict[
       f"""{prefix}.encoder.layer.{teacher_idx}.intermediate.dense.{w}"""
   ]
   lowerCamelCase							:   Dict  						=	state_dict[
       f"""{prefix}.encoder.layer.{teacher_idx}.output.dense.{w}"""
   ]
   lowerCamelCase							:   Any  						=	state_dict[
       f"""{prefix}.encoder.layer.{teacher_idx}.output.LayerNorm.{w}"""
   ]
  std_idx += 1

 lowerCamelCase							:   Optional[int]  						=	state_dict['cls.predictions.decoder.weight']
 lowerCamelCase							:   str  						=	state_dict['cls.predictions.bias']
 if args.vocab_transform:
  for w in ["weight", "bias"]:
   lowerCamelCase							:   str  						=	state_dict[f"""cls.predictions.transform.dense.{w}"""]
   lowerCamelCase							:   Any  						=	state_dict[f"""cls.predictions.transform.LayerNorm.{w}"""]

 print(f"""N layers selected for distillation: {std_idx}""")
 print(f"""Number of params transferred for distillation: {len(compressed_sd.keys())}""")

 print(f"""Save transferred checkpoint to {args.dump_checkpoint}.""")
 torch.save(compressed_sd, args.dump_checkpoint)




