


import gc
import unittest

import numpy as np
import torch
from torch.backends.cuda import sdp_kernel

from diffusers import (
    CMStochasticIterativeScheduler,
    ConsistencyModelPipeline,
    UNetaDModel,
)
from diffusers.utils import randn_tensor, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_a, require_torch_gpu

from ..pipeline_params import UNCONDITIONAL_IMAGE_GENERATION_BATCH_PARAMS, UNCONDITIONAL_IMAGE_GENERATION_PARAMS
from ..test_pipelines_common import PipelineTesterMixin


enable_full_determinism()






class    UpperCAmelCase__	( A_       ,  unittest.TestCase		):



    """simple docstring"""






    UpperCAmelCase__						:		Optional[int]              =  ConsistencyModelPipeline
    UpperCAmelCase__						:		Union[str, Any]              =  UNCONDITIONAL_IMAGE_GENERATION_PARAMS
    UpperCAmelCase__						:		str              =  UNCONDITIONAL_IMAGE_GENERATION_BATCH_PARAMS

    # Override required_optional_params to remove num_images_per_prompt
    UpperCAmelCase__						:		List[str]              =  frozenset(
        [
            "num_inference_steps",
            "generator",
            "latents",
            "output_type",
            "return_dict",
            "callback",
            "callback_steps",
        ]		)






    @property
    def _a (      self   )  ->	Any:
         __UpperCamelCase   							=UNetaDModel.from_pretrained(
             'diffusers/consistency-models-test' ,       subfolder='test_unet' ,       )
         return unet






    @property
    def _a (      self   )  ->	Any:
         __UpperCamelCase   							=UNetaDModel.from_pretrained(
             'diffusers/consistency-models-test' ,       subfolder='test_unet_class_cond' ,       )
         return unet






    def _a (      self ,       A_=False   )  ->	Tuple:
         if class_cond:
              __UpperCamelCase   							=self.dummy_cond_unet
         else:
              __UpperCamelCase   							=self.dummy_uncond_unet

         # Default to CM multistep sampler
         __UpperCamelCase   							=CMStochasticIterativeScheduler(
             num_train_timesteps=40 ,       sigma_min=0.002 ,       sigma_max=80.0 ,       )

         __UpperCamelCase   							={
             'unet': unet,
             'scheduler': scheduler,
         }

         return components






    def _a (      self ,       A_ ,       A_=0   )  ->	List[Any]:
         if str(A_   ).startswith('mps'   ):
              __UpperCamelCase   							=torch.manual_seed(A_   )
         else:
              __UpperCamelCase   							=torch.Generator(device=A_   ).manual_seed(A_   )

         __UpperCamelCase   							={
             'batch_size': 1,
             'num_inference_steps': None,
             'timesteps': [22, 0],
             'generator': generator,
             'output_type': 'np',
         }

         return inputs






    def _a (      self   )  ->	str:
         __UpperCamelCase   							='cpu'  # ensure determinism for the device-dependent torch.Generator
         __UpperCamelCase   							=self.get_dummy_components()
         __UpperCamelCase   							=ConsistencyModelPipeline(**A_   )
         __UpperCamelCase   							=pipe.to(A_   )
         pipe.set_progress_bar_config(disable=A_   )

         __UpperCamelCase   							=self.get_dummy_inputs(A_   )
         __UpperCamelCase   							=pipe(**A_   ).images
         assert image.shape == (1, 32, 32, 3)

         __UpperCamelCase   							=image[0, -3:, -3:, -1]
         __UpperCamelCase   							=np.array([0.3572, 0.6273, 0.4031, 0.3961, 0.4321, 0.5730, 0.5266, 0.4780, 0.5004]   )

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






    def _a (      self   )  ->	Optional[Any]:
         __UpperCamelCase   							='cpu'  # ensure determinism for the device-dependent torch.Generator
         __UpperCamelCase   							=self.get_dummy_components(class_cond=A_   )
         __UpperCamelCase   							=ConsistencyModelPipeline(**A_   )
         __UpperCamelCase   							=pipe.to(A_   )
         pipe.set_progress_bar_config(disable=A_   )

         __UpperCamelCase   							=self.get_dummy_inputs(A_   )
         __UpperCamelCase   							=0
         __UpperCamelCase   							=pipe(**A_   ).images
         assert image.shape == (1, 32, 32, 3)

         __UpperCamelCase   							=image[0, -3:, -3:, -1]
         __UpperCamelCase   							=np.array([0.3572, 0.6273, 0.4031, 0.3961, 0.4321, 0.5730, 0.5266, 0.4780, 0.5004]   )

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






    def _a (      self   )  ->	Dict:
         __UpperCamelCase   							='cpu'  # ensure determinism for the device-dependent torch.Generator
         __UpperCamelCase   							=self.get_dummy_components()
         __UpperCamelCase   							=ConsistencyModelPipeline(**A_   )
         __UpperCamelCase   							=pipe.to(A_   )
         pipe.set_progress_bar_config(disable=A_   )

         __UpperCamelCase   							=self.get_dummy_inputs(A_   )
         __UpperCamelCase   							=1
         __UpperCamelCase   							=None
         __UpperCamelCase   							=pipe(**A_   ).images
         assert image.shape == (1, 32, 32, 3)

         __UpperCamelCase   							=image[0, -3:, -3:, -1]
         __UpperCamelCase   							=np.array([0.5004, 0.5004, 0.4994, 0.5008, 0.4976, 0.5018, 0.4990, 0.4982, 0.4987]   )

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






    def _a (      self   )  ->	int:
         __UpperCamelCase   							='cpu'  # ensure determinism for the device-dependent torch.Generator
         __UpperCamelCase   							=self.get_dummy_components(class_cond=A_   )
         __UpperCamelCase   							=ConsistencyModelPipeline(**A_   )
         __UpperCamelCase   							=pipe.to(A_   )
         pipe.set_progress_bar_config(disable=A_   )

         __UpperCamelCase   							=self.get_dummy_inputs(A_   )
         __UpperCamelCase   							=1
         __UpperCamelCase   							=None
         __UpperCamelCase   							=0
         __UpperCamelCase   							=pipe(**A_   ).images
         assert image.shape == (1, 32, 32, 3)

         __UpperCamelCase   							=image[0, -3:, -3:, -1]
         __UpperCamelCase   							=np.array([0.5004, 0.5004, 0.4994, 0.5008, 0.4976, 0.5018, 0.4990, 0.4982, 0.4987]   )

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






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



    """simple docstring"""






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






    def _a (      self ,       A_=0 ,       A_=False ,       A_="cpu" ,       A_=torch.floataa ,       A_=(1, 3, 64, 64)   )  ->	List[Any]:
         __UpperCamelCase   							=torch.manual_seed(A_   )

         __UpperCamelCase   							={
             'num_inference_steps': None,
             'timesteps': [22, 0],
             'class_labels': 0,
             'generator': generator,
             'output_type': 'np',
         }

         if get_fixed_latents:
              __UpperCamelCase   							=self.get_fixed_latents(seed=A_ ,       device=A_ ,       dtype=A_ ,       shape=A_   )
              __UpperCamelCase   							=latents

         return inputs






    def _a (      self ,       A_=0 ,       A_="cpu" ,       A_=torch.floataa ,       A_=(1, 3, 64, 64)   )  ->	Tuple:
         if type(A_   ) == str:
              __UpperCamelCase   							=torch.device(A_   )
         __UpperCamelCase   							=torch.Generator(device=A_   ).manual_seed(A_   )
         __UpperCamelCase   							=randn_tensor(A_ ,       generator=A_ ,       device=A_ ,       dtype=A_   )
         return latents






    def _a (      self   )  ->	int:
         __UpperCamelCase   							=UNetaDModel.from_pretrained('diffusers/consistency_models' ,       subfolder='diffusers_cd_imagenet64_l2'   )
         __UpperCamelCase   							=CMStochasticIterativeScheduler(
             num_train_timesteps=40 ,       sigma_min=0.002 ,       sigma_max=80.0 ,       )
         __UpperCamelCase   							=ConsistencyModelPipeline(unet=A_ ,       scheduler=A_   )
         pipe.to(torch_device=A_   )
         pipe.set_progress_bar_config(disable=A_   )

         __UpperCamelCase   							=self.get_inputs()
         __UpperCamelCase   							=pipe(**A_   ).images
         assert image.shape == (1, 64, 64, 3)

         __UpperCamelCase   							=image[0, -3:, -3:, -1]

         __UpperCamelCase   							=np.array([0.0888, 0.0881, 0.0666, 0.0479, 0.0292, 0.0195, 0.0201, 0.0163, 0.0254]   )

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






    def _a (      self   )  ->	int:
         __UpperCamelCase   							=UNetaDModel.from_pretrained('diffusers/consistency_models' ,       subfolder='diffusers_cd_imagenet64_l2'   )
         __UpperCamelCase   							=CMStochasticIterativeScheduler(
             num_train_timesteps=40 ,       sigma_min=0.002 ,       sigma_max=80.0 ,       )
         __UpperCamelCase   							=ConsistencyModelPipeline(unet=A_ ,       scheduler=A_   )
         pipe.to(torch_device=A_   )
         pipe.set_progress_bar_config(disable=A_   )

         __UpperCamelCase   							=self.get_inputs()
         __UpperCamelCase   							=1
         __UpperCamelCase   							=None
         __UpperCamelCase   							=pipe(**A_   ).images
         assert image.shape == (1, 64, 64, 3)

         __UpperCamelCase   							=image[0, -3:, -3:, -1]

         __UpperCamelCase   							=np.array([0.0340, 0.0152, 0.0063, 0.0267, 0.0221, 0.0107, 0.0416, 0.0186, 0.0217]   )

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






    @require_torch_a
    def _a (      self   )  ->	Union[str, Any]:
         __UpperCamelCase   							=UNetaDModel.from_pretrained('diffusers/consistency_models' ,       subfolder='diffusers_cd_imagenet64_l2'   )
         __UpperCamelCase   							=CMStochasticIterativeScheduler(
             num_train_timesteps=40 ,       sigma_min=0.002 ,       sigma_max=80.0 ,       )
         __UpperCamelCase   							=ConsistencyModelPipeline(unet=A_ ,       scheduler=A_   )
         pipe.to(torch_device=A_ ,       torch_dtype=torch.floataa   )
         pipe.set_progress_bar_config(disable=A_   )

         __UpperCamelCase   							=self.get_inputs(get_fixed_latents=A_ ,       device=A_   )
         # Ensure usage of flash attention in torch 2.0
         with sdp_kernel(enable_flash=A_ ,       enable_math=A_ ,       enable_mem_efficient=A_   ):
              __UpperCamelCase   							=pipe(**A_   ).images
         assert image.shape == (1, 64, 64, 3)

         __UpperCamelCase   							=image[0, -3:, -3:, -1]

         __UpperCamelCase   							=np.array([0.1875, 0.1428, 0.1289, 0.2151, 0.2092, 0.1477, 0.1877, 0.1641, 0.1353]   )

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






    @require_torch_a
    def _a (      self   )  ->	Union[str, Any]:
         __UpperCamelCase   							=UNetaDModel.from_pretrained('diffusers/consistency_models' ,       subfolder='diffusers_cd_imagenet64_l2'   )
         __UpperCamelCase   							=CMStochasticIterativeScheduler(
             num_train_timesteps=40 ,       sigma_min=0.002 ,       sigma_max=80.0 ,       )
         __UpperCamelCase   							=ConsistencyModelPipeline(unet=A_ ,       scheduler=A_   )
         pipe.to(torch_device=A_ ,       torch_dtype=torch.floataa   )
         pipe.set_progress_bar_config(disable=A_   )

         __UpperCamelCase   							=self.get_inputs(get_fixed_latents=A_ ,       device=A_   )
         __UpperCamelCase   							=1
         __UpperCamelCase   							=None
         # Ensure usage of flash attention in torch 2.0
         with sdp_kernel(enable_flash=A_ ,       enable_math=A_ ,       enable_mem_efficient=A_   ):
              __UpperCamelCase   							=pipe(**A_   ).images
         assert image.shape == (1, 64, 64, 3)

         __UpperCamelCase   							=image[0, -3:, -3:, -1]

         __UpperCamelCase   							=np.array([0.1663, 0.1948, 0.2275, 0.1680, 0.1204, 0.1245, 0.1858, 0.1338, 0.2095]   )

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









import argparse

import torch

from transformers import BertForMaskedLM


if __name__ == "__main__":
   _A			=						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')
   _A			=						parser.parse_args()

   if args.model_type == "bert":
      _A			=						BertForMaskedLM.from_pretrained(args.model_name)
      _A			=						'bert'
   else:
      raise ValueError('args.model_type should be "bert".')

   _A			=						model.state_dict()
   _A			=						{}

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

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

         _A			=						state_dict[
             f"""{prefix}.encoder.layer.{teacher_idx}.attention.output.dense.{w}"""
         ]
         _A			=						state_dict[
             f"""{prefix}.encoder.layer.{teacher_idx}.attention.output.LayerNorm.{w}"""
         ]

         _A			=						state_dict[
             f"""{prefix}.encoder.layer.{teacher_idx}.intermediate.dense.{w}"""
         ]
         _A			=						state_dict[
             f"""{prefix}.encoder.layer.{teacher_idx}.output.dense.{w}"""
         ]
         _A			=						state_dict[
             f"""{prefix}.encoder.layer.{teacher_idx}.output.LayerNorm.{w}"""
         ]
      std_idx += 1

   _A			=						state_dict['cls.predictions.decoder.weight']
   _A			=						state_dict['cls.predictions.bias']
   if args.vocab_transform:
      for w in ["weight", "bias"]:
         _A			=						state_dict[f"""cls.predictions.transform.dense.{w}"""]
         _A			=						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)





