






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   lowerCamelCase      (_snake_case ,  unittest.TestCase			):







		'''simple docstring'''



		_snake_case    :    Tuple       						=  ConsistencyModelPipeline
		_snake_case    :    Union[str, Any]       						=  UNCONDITIONAL_IMAGE_GENERATION_PARAMS
		_snake_case    :    Union[str, Any]       						=  UNCONDITIONAL_IMAGE_GENERATION_BATCH_PARAMS

		# Override required_optional_params to remove num_images_per_prompt
		_snake_case    :    Optional[Any]       						=  frozenset(
		    [
		        '''num_inference_steps''',
		        '''generator''',
		        '''latents''',
		        '''output_type''',
		        '''return_dict''',
		        '''callback''',
		        '''callback_steps''',
		    ]			)



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



		@property
		def 					__UpperCAmelCase			(							self   )	->							Optional[int]:
							UpperCAmelCase_    :      str            =						UNetaDModel.from_pretrained(
							    'diffusers/consistency-models-test'				,		subfolder='test_unet_class_cond'				,		)
							return unet



		def 					__UpperCAmelCase			(							self				,		_UpperCamelCase=False   )	->							Any:
							if class_cond:
												UpperCAmelCase_    :      List[Any]            =						self.dummy_cond_unet
							else:
												UpperCAmelCase_    :      int            =						self.dummy_uncond_unet

							# Default to CM multistep sampler
							UpperCAmelCase_    :      Dict            =						CMStochasticIterativeScheduler(
							    num_train_timesteps=4_0				,		sigma_min=0.0_02				,		sigma_max=80.0				,		)

							UpperCAmelCase_    :      int            =						{
							    'unet': unet,
							    'scheduler': scheduler,
							}

							return components



		def 					__UpperCAmelCase			(							self				,		_UpperCamelCase				,		_UpperCamelCase=0   )	->							List[Any]:
							if str(_UpperCamelCase   ).startswith('mps'   ):
												UpperCAmelCase_    :      Optional[Any]            =						torch.manual_seed(_UpperCamelCase   )
							else:
												UpperCAmelCase_    :      Dict            =						torch.Generator(device=_UpperCamelCase   ).manual_seed(_UpperCamelCase   )

							UpperCAmelCase_    :      Tuple            =						{
							    'batch_size': 1,
							    'num_inference_steps': None,
							    'timesteps': [2_2, 0],
							    'generator': generator,
							    'output_type': 'np',
							}

							return inputs



		def 					__UpperCAmelCase			(							self   )	->							Tuple:
							UpperCAmelCase_    :      Any            =						'cpu'  # ensure determinism for the device-dependent torch.Generator
							UpperCAmelCase_    :      Optional[Any]            =						self.get_dummy_components()
							UpperCAmelCase_    :      Dict            =						ConsistencyModelPipeline(**_UpperCamelCase   )
							UpperCAmelCase_    :      List[Any]            =						pipe.to(_UpperCamelCase   )
							pipe.set_progress_bar_config(disable=_UpperCamelCase   )

							UpperCAmelCase_    :      Optional[int]            =						self.get_dummy_inputs(_UpperCamelCase   )
							UpperCAmelCase_    :      Optional[Any]            =						pipe(**_UpperCamelCase   ).images
							assert image.shape == (1, 3_2, 3_2, 3)

							UpperCAmelCase_    :      Optional[int]            =						image[0, -3:, -3:, -1]
							UpperCAmelCase_    :      List[str]            =						np.array([0.35_72, 0.62_73, 0.40_31, 0.39_61, 0.43_21, 0.57_30, 0.52_66, 0.47_80, 0.50_04]   )

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



		def 					__UpperCAmelCase			(							self   )	->							List[str]:
							UpperCAmelCase_    :      Dict            =						'cpu'  # ensure determinism for the device-dependent torch.Generator
							UpperCAmelCase_    :      str            =						self.get_dummy_components(class_cond=_UpperCamelCase   )
							UpperCAmelCase_    :      List[Any]            =						ConsistencyModelPipeline(**_UpperCamelCase   )
							UpperCAmelCase_    :      Optional[Any]            =						pipe.to(_UpperCamelCase   )
							pipe.set_progress_bar_config(disable=_UpperCamelCase   )

							UpperCAmelCase_    :      List[Any]            =						self.get_dummy_inputs(_UpperCamelCase   )
							UpperCAmelCase_    :      int            =						0
							UpperCAmelCase_    :      Dict            =						pipe(**_UpperCamelCase   ).images
							assert image.shape == (1, 3_2, 3_2, 3)

							UpperCAmelCase_    :      Dict            =						image[0, -3:, -3:, -1]
							UpperCAmelCase_    :      List[str]            =						np.array([0.35_72, 0.62_73, 0.40_31, 0.39_61, 0.43_21, 0.57_30, 0.52_66, 0.47_80, 0.50_04]   )

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



		def 					__UpperCAmelCase			(							self   )	->							Dict:
							UpperCAmelCase_    :      Tuple            =						'cpu'  # ensure determinism for the device-dependent torch.Generator
							UpperCAmelCase_    :      List[Any]            =						self.get_dummy_components()
							UpperCAmelCase_    :      Union[str, Any]            =						ConsistencyModelPipeline(**_UpperCamelCase   )
							UpperCAmelCase_    :      Union[str, Any]            =						pipe.to(_UpperCamelCase   )
							pipe.set_progress_bar_config(disable=_UpperCamelCase   )

							UpperCAmelCase_    :      List[Any]            =						self.get_dummy_inputs(_UpperCamelCase   )
							UpperCAmelCase_    :      int            =						1
							UpperCAmelCase_    :      Tuple            =						None
							UpperCAmelCase_    :      str            =						pipe(**_UpperCamelCase   ).images
							assert image.shape == (1, 3_2, 3_2, 3)

							UpperCAmelCase_    :      Optional[Any]            =						image[0, -3:, -3:, -1]
							UpperCAmelCase_    :      Optional[int]            =						np.array([0.50_04, 0.50_04, 0.49_94, 0.50_08, 0.49_76, 0.50_18, 0.49_90, 0.49_82, 0.49_87]   )

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



		def 					__UpperCAmelCase			(							self   )	->							Dict:
							UpperCAmelCase_    :      int            =						'cpu'  # ensure determinism for the device-dependent torch.Generator
							UpperCAmelCase_    :      Tuple            =						self.get_dummy_components(class_cond=_UpperCamelCase   )
							UpperCAmelCase_    :      Union[str, Any]            =						ConsistencyModelPipeline(**_UpperCamelCase   )
							UpperCAmelCase_    :      List[Any]            =						pipe.to(_UpperCamelCase   )
							pipe.set_progress_bar_config(disable=_UpperCamelCase   )

							UpperCAmelCase_    :      str            =						self.get_dummy_inputs(_UpperCamelCase   )
							UpperCAmelCase_    :      Tuple            =						1
							UpperCAmelCase_    :      Optional[int]            =						None
							UpperCAmelCase_    :      Optional[int]            =						0
							UpperCAmelCase_    :      List[str]            =						pipe(**_UpperCamelCase   ).images
							assert image.shape == (1, 3_2, 3_2, 3)

							UpperCAmelCase_    :      str            =						image[0, -3:, -3:, -1]
							UpperCAmelCase_    :      Tuple            =						np.array([0.50_04, 0.50_04, 0.49_94, 0.50_08, 0.49_76, 0.50_18, 0.49_90, 0.49_82, 0.49_87]   )

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





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







		'''simple docstring'''



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



		def 					__UpperCAmelCase			(							self				,		_UpperCamelCase=0				,		_UpperCamelCase=False				,		_UpperCamelCase="cpu"				,		_UpperCamelCase=torch.floataa				,		_UpperCamelCase=(1, 3, 6_4, 6_4)   )	->							int:
							UpperCAmelCase_    :      List[str]            =						torch.manual_seed(_UpperCamelCase   )

							UpperCAmelCase_    :      Optional[int]            =						{
							    'num_inference_steps': None,
							    'timesteps': [2_2, 0],
							    'class_labels': 0,
							    'generator': generator,
							    'output_type': 'np',
							}

							if get_fixed_latents:
												UpperCAmelCase_    :      Any            =						self.get_fixed_latents(seed=_UpperCamelCase				,		device=_UpperCamelCase				,		dtype=_UpperCamelCase				,		shape=_UpperCamelCase   )
												UpperCAmelCase_    :      Optional[Any]            =						latents

							return inputs



		def 					__UpperCAmelCase			(							self				,		_UpperCamelCase=0				,		_UpperCamelCase="cpu"				,		_UpperCamelCase=torch.floataa				,		_UpperCamelCase=(1, 3, 6_4, 6_4)   )	->							Tuple:
							if type(_UpperCamelCase   ) == str:
												UpperCAmelCase_    :      str            =						torch.device(_UpperCamelCase   )
							UpperCAmelCase_    :      Union[str, Any]            =						torch.Generator(device=_UpperCamelCase   ).manual_seed(_UpperCamelCase   )
							UpperCAmelCase_    :      Optional[Any]            =						randn_tensor(_UpperCamelCase				,		generator=_UpperCamelCase				,		device=_UpperCamelCase				,		dtype=_UpperCamelCase   )
							return latents



		def 					__UpperCAmelCase			(							self   )	->							Any:
							UpperCAmelCase_    :      int            =						UNetaDModel.from_pretrained('diffusers/consistency_models'				,		subfolder='diffusers_cd_imagenet64_l2'   )
							UpperCAmelCase_    :      Optional[Any]            =						CMStochasticIterativeScheduler(
							    num_train_timesteps=4_0				,		sigma_min=0.0_02				,		sigma_max=80.0				,		)
							UpperCAmelCase_    :      List[Any]            =						ConsistencyModelPipeline(unet=_UpperCamelCase				,		scheduler=_UpperCamelCase   )
							pipe.to(torch_device=_UpperCamelCase   )
							pipe.set_progress_bar_config(disable=_UpperCamelCase   )

							UpperCAmelCase_    :      List[Any]            =						self.get_inputs()
							UpperCAmelCase_    :      Any            =						pipe(**_UpperCamelCase   ).images
							assert image.shape == (1, 6_4, 6_4, 3)

							UpperCAmelCase_    :      str            =						image[0, -3:, -3:, -1]

							UpperCAmelCase_    :      Optional[int]            =						np.array([0.08_88, 0.08_81, 0.06_66, 0.04_79, 0.02_92, 0.01_95, 0.02_01, 0.01_63, 0.02_54]   )

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



		def 					__UpperCAmelCase			(							self   )	->							int:
							UpperCAmelCase_    :      Dict            =						UNetaDModel.from_pretrained('diffusers/consistency_models'				,		subfolder='diffusers_cd_imagenet64_l2'   )
							UpperCAmelCase_    :      str            =						CMStochasticIterativeScheduler(
							    num_train_timesteps=4_0				,		sigma_min=0.0_02				,		sigma_max=80.0				,		)
							UpperCAmelCase_    :      Optional[Any]            =						ConsistencyModelPipeline(unet=_UpperCamelCase				,		scheduler=_UpperCamelCase   )
							pipe.to(torch_device=_UpperCamelCase   )
							pipe.set_progress_bar_config(disable=_UpperCamelCase   )

							UpperCAmelCase_    :      Optional[int]            =						self.get_inputs()
							UpperCAmelCase_    :      int            =						1
							UpperCAmelCase_    :      List[str]            =						None
							UpperCAmelCase_    :      int            =						pipe(**_UpperCamelCase   ).images
							assert image.shape == (1, 6_4, 6_4, 3)

							UpperCAmelCase_    :      Optional[int]            =						image[0, -3:, -3:, -1]

							UpperCAmelCase_    :      Optional[int]            =						np.array([0.03_40, 0.01_52, 0.00_63, 0.02_67, 0.02_21, 0.01_07, 0.04_16, 0.01_86, 0.02_17]   )

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



		@require_torch_a
		def 					__UpperCAmelCase			(							self   )	->							List[Any]:
							UpperCAmelCase_    :      List[str]            =						UNetaDModel.from_pretrained('diffusers/consistency_models'				,		subfolder='diffusers_cd_imagenet64_l2'   )
							UpperCAmelCase_    :      int            =						CMStochasticIterativeScheduler(
							    num_train_timesteps=4_0				,		sigma_min=0.0_02				,		sigma_max=80.0				,		)
							UpperCAmelCase_    :      Union[str, Any]            =						ConsistencyModelPipeline(unet=_UpperCamelCase				,		scheduler=_UpperCamelCase   )
							pipe.to(torch_device=_UpperCamelCase				,		torch_dtype=torch.floataa   )
							pipe.set_progress_bar_config(disable=_UpperCamelCase   )

							UpperCAmelCase_    :      Optional[Any]            =						self.get_inputs(get_fixed_latents=_UpperCamelCase				,		device=_UpperCamelCase   )
							# Ensure usage of flash attention in torch 2.0
							with sdp_kernel(enable_flash=_UpperCamelCase				,		enable_math=_UpperCamelCase				,		enable_mem_efficient=_UpperCamelCase   ):
												UpperCAmelCase_    :      List[Any]            =						pipe(**_UpperCamelCase   ).images
							assert image.shape == (1, 6_4, 6_4, 3)

							UpperCAmelCase_    :      Any            =						image[0, -3:, -3:, -1]

							UpperCAmelCase_    :      Tuple            =						np.array([0.18_75, 0.14_28, 0.12_89, 0.21_51, 0.20_92, 0.14_77, 0.18_77, 0.16_41, 0.13_53]   )

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



		@require_torch_a
		def 					__UpperCAmelCase			(							self   )	->							List[str]:
							UpperCAmelCase_    :      List[Any]            =						UNetaDModel.from_pretrained('diffusers/consistency_models'				,		subfolder='diffusers_cd_imagenet64_l2'   )
							UpperCAmelCase_    :      List[Any]            =						CMStochasticIterativeScheduler(
							    num_train_timesteps=4_0				,		sigma_min=0.0_02				,		sigma_max=80.0				,		)
							UpperCAmelCase_    :      int            =						ConsistencyModelPipeline(unet=_UpperCamelCase				,		scheduler=_UpperCamelCase   )
							pipe.to(torch_device=_UpperCamelCase				,		torch_dtype=torch.floataa   )
							pipe.set_progress_bar_config(disable=_UpperCamelCase   )

							UpperCAmelCase_    :      Dict            =						self.get_inputs(get_fixed_latents=_UpperCamelCase				,		device=_UpperCamelCase   )
							UpperCAmelCase_    :      List[str]            =						1
							UpperCAmelCase_    :      str            =						None
							# Ensure usage of flash attention in torch 2.0
							with sdp_kernel(enable_flash=_UpperCamelCase				,		enable_math=_UpperCamelCase				,		enable_mem_efficient=_UpperCamelCase   ):
												UpperCAmelCase_    :      Tuple            =						pipe(**_UpperCamelCase   ).images
							assert image.shape == (1, 6_4, 6_4, 3)

							UpperCAmelCase_    :      Dict            =						image[0, -3:, -3:, -1]

							UpperCAmelCase_    :      str            =						np.array([0.16_63, 0.19_48, 0.22_75, 0.16_80, 0.12_04, 0.12_45, 0.18_58, 0.13_38, 0.20_95]   )

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







import argparse

import torch

from transformers import GPTaLMHeadModel, RobertaForMaskedLM


if __name__ == "__main__":
   a_    					=       argparse.ArgumentParser(
       description=(
           'Extraction some layers of the full RobertaForMaskedLM or GPT2LMHeadModel for Transfer Learned'
           ' Distillation'
       )
   )
   parser.add_argument('--model_type', default='roberta', choices=['roberta', 'gpt2'])
   parser.add_argument('--model_name', default='roberta-large', type=str)
   parser.add_argument('--dump_checkpoint', default='serialization_dir/tf_roberta_048131723.pth', type=str)
   parser.add_argument('--vocab_transform', action='store_true')
   a_    					=       parser.parse_args()

   if args.model_type == "roberta":
      a_    					=       RobertaForMaskedLM.from_pretrained(args.model_name)
      a_    					=       'roberta'
   elif args.model_type == "gpt2":
      a_    					=       GPTaLMHeadModel.from_pretrained(args.model_name)
      a_    					=       'transformer'

   a_    					=       model.state_dict()
   a_    					=       {}

   # Embeddings #
   if args.model_type == "gpt2":
      for param_name in ["wte.weight", "wpe.weight"]:
         a_    					=       state_dict[F"""{prefix}.{param_name}"""]
   else:
      for w in ["word_embeddings", "position_embeddings", "token_type_embeddings"]:
         a_    					=       F"""{prefix}.embeddings.{w}.weight"""
         a_    					=       state_dict[param_name]
      for w in ["weight", "bias"]:
         a_    					=       F"""{prefix}.embeddings.LayerNorm.{w}"""
         a_    					=       state_dict[param_name]

    # Transformer Blocks #
   a_    					=       0
   for teacher_idx in [0, 2, 4, 7, 9, 11]:
      if args.model_type == "gpt2":
         for layer in ["ln_1", "attn.c_attn", "attn.c_proj", "ln_2", "mlp.c_fc", "mlp.c_proj"]:
            for w in ["weight", "bias"]:
               a_    					=       state_dict[
                   F"""{prefix}.h.{teacher_idx}.{layer}.{w}"""
               ]
         a_    					=       state_dict[F"""{prefix}.h.{teacher_idx}.attn.bias"""]
      else:
         for layer in [
             "attention.self.query",
             "attention.self.key",
             "attention.self.value",
             "attention.output.dense",
             "attention.output.LayerNorm",
             "intermediate.dense",
             "output.dense",
             "output.LayerNorm",
         ]:
            for w in ["weight", "bias"]:
               a_    					=       state_dict[
                   F"""{prefix}.encoder.layer.{teacher_idx}.{layer}.{w}"""
               ]
      std_idx += 1

   # Language Modeling Head ###s
   if args.model_type == "roberta":
      for layer in ["lm_head.decoder.weight", "lm_head.bias"]:
         a_    					=       state_dict[F"""{layer}"""]
      if args.vocab_transform:
         for w in ["weight", "bias"]:
            a_    					=       state_dict[F"""lm_head.dense.{w}"""]
            a_    					=       state_dict[F"""lm_head.layer_norm.{w}"""]
   elif args.model_type == "gpt2":
      for w in ["weight", "bias"]:
         a_    					=       state_dict[F"""{prefix}.ln_f.{w}"""]
      a_    					=       state_dict['lm_head.weight']

   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)



