





"""simple docstring"""





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  (	a_    ,  unittest.TestCase			):
       """simple docstring"""




       __magic_name__	:Optional[int]						    =   ConsistencyModelPipeline
       __magic_name__	:Optional[int]						    =   UNCONDITIONAL_IMAGE_GENERATION_PARAMS
       __magic_name__	:Tuple						    =   UNCONDITIONAL_IMAGE_GENERATION_BATCH_PARAMS

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






       @property
       def 	snake_case		(       self ):



         '''simple docstring'''

         lowerCAmelCase__     :Any											=				UNetaDModel.from_pretrained(
             'diffusers/consistency-models-test'       ,					subfolder='test_unet'       ,					)
         return unet






       @property
       def 	snake_case		(       self ):



         '''simple docstring'''

         lowerCAmelCase__     :int											=				UNetaDModel.from_pretrained(
             'diffusers/consistency-models-test'       ,					subfolder='test_unet_class_cond'       ,					)
         return unet






       def 	snake_case		(       self       ,					__UpperCAmelCase=False ):



         '''simple docstring'''

         if class_cond:
           lowerCAmelCase__     :Tuple											=				self.dummy_cond_unet
         else:
           lowerCAmelCase__     :Any											=				self.dummy_uncond_unet

         # Default to CM multistep sampler
         lowerCAmelCase__     :Optional[int]											=				CMStochasticIterativeScheduler(
             num_train_timesteps=4_0       ,					sigma_min=0.0_02       ,					sigma_max=80.0       ,					)

         lowerCAmelCase__     :List[str]											=				{
             'unet': unet,
             'scheduler': scheduler,
         }

         return components






       def 	snake_case		(       self       ,					__UpperCAmelCase       ,					__UpperCAmelCase=0 ):



         '''simple docstring'''

         if str(_SCREAMING_SNAKE_CASE ).startswith('mps' ):
           lowerCAmelCase__     :Optional[Any]											=				torch.manual_seed(_SCREAMING_SNAKE_CASE )
         else:
           lowerCAmelCase__     :Dict											=				torch.Generator(device=_SCREAMING_SNAKE_CASE ).manual_seed(_SCREAMING_SNAKE_CASE )

         lowerCAmelCase__     :Optional[int]											=				{
             'batch_size': 1,
             'num_inference_steps': None,
             'timesteps': [2_2, 0],
             'generator': generator,
             'output_type': 'np',
         }

         return inputs






       def 	snake_case		(       self ):



         '''simple docstring'''

         lowerCAmelCase__     :Optional[int]											=				'cpu'  # ensure determinism for the device-dependent torch.Generator
         lowerCAmelCase__     :int											=				self.get_dummy_components()
         lowerCAmelCase__     :List[Any]											=				ConsistencyModelPipeline(**_SCREAMING_SNAKE_CASE )
         lowerCAmelCase__     :List[str]											=				pipe.to(_SCREAMING_SNAKE_CASE )
         pipe.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE )

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

         lowerCAmelCase__     :Dict											=				image[0, -3:, -3:, -1]
         lowerCAmelCase__     :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 	snake_case		(       self ):



         '''simple docstring'''

         lowerCAmelCase__     :Tuple											=				'cpu'  # ensure determinism for the device-dependent torch.Generator
         lowerCAmelCase__     :Optional[Any]											=				self.get_dummy_components(class_cond=_SCREAMING_SNAKE_CASE )
         lowerCAmelCase__     :List[Any]											=				ConsistencyModelPipeline(**_SCREAMING_SNAKE_CASE )
         lowerCAmelCase__     :int											=				pipe.to(_SCREAMING_SNAKE_CASE )
         pipe.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE )

         lowerCAmelCase__     :List[str]											=				self.get_dummy_inputs(_SCREAMING_SNAKE_CASE )
         lowerCAmelCase__     :Union[str, Any]											=				0
         lowerCAmelCase__     :List[Any]											=				pipe(**_SCREAMING_SNAKE_CASE ).images
         assert image.shape == (1, 3_2, 3_2, 3)

         lowerCAmelCase__     :Dict											=				image[0, -3:, -3:, -1]
         lowerCAmelCase__     :Tuple											=				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 	snake_case		(       self ):



         '''simple docstring'''

         lowerCAmelCase__     :List[str]											=				'cpu'  # ensure determinism for the device-dependent torch.Generator
         lowerCAmelCase__     :Tuple											=				self.get_dummy_components()
         lowerCAmelCase__     :str											=				ConsistencyModelPipeline(**_SCREAMING_SNAKE_CASE )
         lowerCAmelCase__     :List[str]											=				pipe.to(_SCREAMING_SNAKE_CASE )
         pipe.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE )

         lowerCAmelCase__     :Union[str, Any]											=				self.get_dummy_inputs(_SCREAMING_SNAKE_CASE )
         lowerCAmelCase__     :str											=				1
         lowerCAmelCase__     :str											=				None
         lowerCAmelCase__     :int											=				pipe(**_SCREAMING_SNAKE_CASE ).images
         assert image.shape == (1, 3_2, 3_2, 3)

         lowerCAmelCase__     :Optional[Any]											=				image[0, -3:, -3:, -1]
         lowerCAmelCase__     :Dict											=				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 	snake_case		(       self ):



         '''simple docstring'''

         lowerCAmelCase__     :List[str]											=				'cpu'  # ensure determinism for the device-dependent torch.Generator
         lowerCAmelCase__     :str											=				self.get_dummy_components(class_cond=_SCREAMING_SNAKE_CASE )
         lowerCAmelCase__     :Dict											=				ConsistencyModelPipeline(**_SCREAMING_SNAKE_CASE )
         lowerCAmelCase__     :Union[str, Any]											=				pipe.to(_SCREAMING_SNAKE_CASE )
         pipe.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE )

         lowerCAmelCase__     :Union[str, Any]											=				self.get_dummy_inputs(_SCREAMING_SNAKE_CASE )
         lowerCAmelCase__     :Union[str, Any]											=				1
         lowerCAmelCase__     :Optional[int]											=				None
         lowerCAmelCase__     :List[Any]											=				0
         lowerCAmelCase__     :Any											=				pipe(**_SCREAMING_SNAKE_CASE ).images
         assert image.shape == (1, 3_2, 3_2, 3)

         lowerCAmelCase__     :Dict											=				image[0, -3:, -3:, -1]
         lowerCAmelCase__     :Dict											=				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 	snake_case		(       self ):



         '''simple docstring'''

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






       def 	snake_case		(       self       ,					__UpperCAmelCase=0       ,					__UpperCAmelCase=False       ,					__UpperCAmelCase="cpu"       ,					__UpperCAmelCase=torch.floataa       ,					__UpperCAmelCase=(1, 3, 6_4, 6_4) ):



         '''simple docstring'''

         lowerCAmelCase__     :str											=				torch.manual_seed(_SCREAMING_SNAKE_CASE )

         lowerCAmelCase__     :Union[str, Any]											=				{
             'num_inference_steps': None,
             'timesteps': [2_2, 0],
             'class_labels': 0,
             'generator': generator,
             'output_type': 'np',
         }

         if get_fixed_latents:
           lowerCAmelCase__     :Tuple											=				self.get_fixed_latents(seed=_SCREAMING_SNAKE_CASE       ,					device=_SCREAMING_SNAKE_CASE       ,					dtype=_SCREAMING_SNAKE_CASE       ,					shape=_SCREAMING_SNAKE_CASE )
           lowerCAmelCase__     :Union[str, Any]											=				latents

         return inputs






       def 	snake_case		(       self       ,					__UpperCAmelCase=0       ,					__UpperCAmelCase="cpu"       ,					__UpperCAmelCase=torch.floataa       ,					__UpperCAmelCase=(1, 3, 6_4, 6_4) ):



         '''simple docstring'''

         if type(_SCREAMING_SNAKE_CASE ) == str:
           lowerCAmelCase__     :Optional[int]											=				torch.device(_SCREAMING_SNAKE_CASE )
         lowerCAmelCase__     :Dict											=				torch.Generator(device=_SCREAMING_SNAKE_CASE ).manual_seed(_SCREAMING_SNAKE_CASE )
         lowerCAmelCase__     :Union[str, Any]											=				randn_tensor(_SCREAMING_SNAKE_CASE       ,					generator=_SCREAMING_SNAKE_CASE       ,					device=_SCREAMING_SNAKE_CASE       ,					dtype=_SCREAMING_SNAKE_CASE )
         return latents






       def 	snake_case		(       self ):



         '''simple docstring'''

         lowerCAmelCase__     :str											=				UNetaDModel.from_pretrained('diffusers/consistency_models'       ,					subfolder='diffusers_cd_imagenet64_l2' )
         lowerCAmelCase__     :int											=				CMStochasticIterativeScheduler(
             num_train_timesteps=4_0       ,					sigma_min=0.0_02       ,					sigma_max=80.0       ,					)
         lowerCAmelCase__     :Any											=				ConsistencyModelPipeline(unet=_SCREAMING_SNAKE_CASE       ,					scheduler=_SCREAMING_SNAKE_CASE )
         pipe.to(torch_device=_SCREAMING_SNAKE_CASE )
         pipe.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE )

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

         lowerCAmelCase__     :Tuple											=				image[0, -3:, -3:, -1]

         lowerCAmelCase__     :List[Any]											=				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 	snake_case		(       self ):



         '''simple docstring'''

         lowerCAmelCase__     :int											=				UNetaDModel.from_pretrained('diffusers/consistency_models'       ,					subfolder='diffusers_cd_imagenet64_l2' )
         lowerCAmelCase__     :Union[str, Any]											=				CMStochasticIterativeScheduler(
             num_train_timesteps=4_0       ,					sigma_min=0.0_02       ,					sigma_max=80.0       ,					)
         lowerCAmelCase__     :Any											=				ConsistencyModelPipeline(unet=_SCREAMING_SNAKE_CASE       ,					scheduler=_SCREAMING_SNAKE_CASE )
         pipe.to(torch_device=_SCREAMING_SNAKE_CASE )
         pipe.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE )

         lowerCAmelCase__     :Optional[Any]											=				self.get_inputs()
         lowerCAmelCase__     :Optional[Any]											=				1
         lowerCAmelCase__     :List[Any]											=				None
         lowerCAmelCase__     :Dict											=				pipe(**_SCREAMING_SNAKE_CASE ).images
         assert image.shape == (1, 6_4, 6_4, 3)

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

         lowerCAmelCase__     :Any											=				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 	snake_case		(       self ):



         '''simple docstring'''

         lowerCAmelCase__     :List[Any]											=				UNetaDModel.from_pretrained('diffusers/consistency_models'       ,					subfolder='diffusers_cd_imagenet64_l2' )
         lowerCAmelCase__     :List[Any]											=				CMStochasticIterativeScheduler(
             num_train_timesteps=4_0       ,					sigma_min=0.0_02       ,					sigma_max=80.0       ,					)
         lowerCAmelCase__     :str											=				ConsistencyModelPipeline(unet=_SCREAMING_SNAKE_CASE       ,					scheduler=_SCREAMING_SNAKE_CASE )
         pipe.to(torch_device=_SCREAMING_SNAKE_CASE       ,					torch_dtype=torch.floataa )
         pipe.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE )

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

         lowerCAmelCase__     :Union[str, Any]											=				image[0, -3:, -3:, -1]

         lowerCAmelCase__     :str											=				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 	snake_case		(       self ):



         '''simple docstring'''

         lowerCAmelCase__     :List[Any]											=				UNetaDModel.from_pretrained('diffusers/consistency_models'       ,					subfolder='diffusers_cd_imagenet64_l2' )
         lowerCAmelCase__     :Optional[int]											=				CMStochasticIterativeScheduler(
             num_train_timesteps=4_0       ,					sigma_min=0.0_02       ,					sigma_max=80.0       ,					)
         lowerCAmelCase__     :List[Any]											=				ConsistencyModelPipeline(unet=_SCREAMING_SNAKE_CASE       ,					scheduler=_SCREAMING_SNAKE_CASE )
         pipe.to(torch_device=_SCREAMING_SNAKE_CASE       ,					torch_dtype=torch.floataa )
         pipe.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE )

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

         lowerCAmelCase__     :Union[str, Any]											=				image[0, -3:, -3:, -1]

         lowerCAmelCase__     :int											=				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



'''simple docstring'''



import argparse
import logging
from collections import namedtuple

import torch
from model_bertabs import BertAbsSummarizer
from models.model_builder import AbsSummarizer  # The authors' implementation

from transformers import BertTokenizer


logging.basicConfig(level=logging.INFO)
_UpperCamelCase						:  Optional[Any]							       =				logging.getLogger(__name__)


_UpperCamelCase						:  List[Any]							       =				"Hello world! cécé herlolip"


_UpperCamelCase						:  int							       =				namedtuple(
    "BertAbsConfig",
    [
        "temp_dir",
        "large",
        "use_bert_emb",
        "finetune_bert",
        "encoder",
        "share_emb",
        "max_pos",
        "enc_layers",
        "enc_hidden_size",
        "enc_heads",
        "enc_ff_size",
        "enc_dropout",
        "dec_layers",
        "dec_hidden_size",
        "dec_heads",
        "dec_ff_size",
        "dec_dropout",
    ],
)


def       snake_case   (	snake_case      : int					,		snake_case      : Optional[Any]			) ->  Optional[Any]:


    """simple docstring"""


    lowerCAmelCase 						=							BertAbsConfig(
        temp_dir='.'					,		finetune_bert=snake_case					,		large=snake_case					,		share_emb=snake_case					,		use_bert_emb=snake_case					,		encoder='bert'					,		max_pos=512					,		enc_layers=6					,		enc_hidden_size=512					,		enc_heads=8					,		enc_ff_size=512					,		enc_dropout=0.2					,		dec_layers=6					,		dec_hidden_size=768					,		dec_heads=8					,		dec_ff_size=2048					,		dec_dropout=0.2					,		)
    lowerCAmelCase 						=							torch.load(snake_case					,		lambda snake_case					,		snake_case			: storage			)
    lowerCAmelCase 						=							AbsSummarizer(snake_case					,		torch.device('cpu'			)					,		snake_case			)
    original.eval()

    lowerCAmelCase 						=							BertAbsSummarizer(snake_case					,		torch.device('cpu'			)			)
    new_model.eval()

    # -------------------
    # Convert the weights
    # -------------------

    logging.info('convert the model'			)
    new_model.bert.load_state_dict(original.bert.state_dict()			)
    new_model.decoder.load_state_dict(original.decoder.state_dict()			)
    new_model.generator.load_state_dict(original.generator.state_dict()			)

    # ----------------------------------
    # Make sure the outpus are identical
    # ----------------------------------

    logging.info('Make sure that the models\' outputs are identical'			)
    lowerCAmelCase 						=							BertTokenizer.from_pretrained('bert-base-uncased'			)

    # prepare the model inputs
    lowerCAmelCase 						=							tokenizer.encode('This is sample éàalj\'-.'			)
    encoder_input_ids.extend([tokenizer.pad_token_id] * (512 - len(snake_case			))			)
    lowerCAmelCase 						=							torch.tensor(snake_case			).unsqueeze(0			)
    lowerCAmelCase 						=							tokenizer.encode('This is sample 3 éàalj\'-.'			)
    decoder_input_ids.extend([tokenizer.pad_token_id] * (512 - len(snake_case			))			)
    lowerCAmelCase 						=							torch.tensor(snake_case			).unsqueeze(0			)

    # failsafe to make sure the weights reset does not affect the
    # loaded weights.
    assert torch.max(torch.abs(original.generator[0].weight - new_model.generator[0].weight			)			) == 0

    # forward pass
    lowerCAmelCase 						=							encoder_input_ids
    lowerCAmelCase 						=							decoder_input_ids
    lowerCAmelCase 						=							lowerCAmelCase 						=							None
    lowerCAmelCase 						=							None
    lowerCAmelCase 						=							lowerCAmelCase 						=							None
    lowerCAmelCase 						=							lowerCAmelCase 						=							None
    lowerCAmelCase 						=							None

    # The original model does not apply the geneator layer immediatly but rather in
    # the beam search (where it combines softmax + linear layer). Since we already
    # apply the softmax in our generation process we only apply the linear layer here.
    # We make sure that the outputs of the full stack are identical
    lowerCAmelCase 						=							original(snake_case					,		snake_case					,		snake_case					,		snake_case					,		snake_case					,		snake_case					,		snake_case			)[0]
    lowerCAmelCase 						=							original.generator(snake_case			)

    lowerCAmelCase 						=							new_model(
        snake_case					,		snake_case					,		snake_case					,		snake_case					,		snake_case			)[0]
    lowerCAmelCase 						=							new_model.generator(snake_case			)

    lowerCAmelCase 						=							torch.max(torch.abs(output_converted_model - output_original_model			)			).item()
    print('Maximum absolute difference beween weights: {:.2f}'.format(snake_case			)			)
    lowerCAmelCase 						=							torch.max(torch.abs(output_converted_generator - output_original_generator			)			).item()
    print('Maximum absolute difference beween weights: {:.2f}'.format(snake_case			)			)

    lowerCAmelCase 						=							torch.allclose(snake_case					,		snake_case					,		atol=1e-3			)
    if are_identical:
        logging.info('all weights are equal up to 1e-3'			)
    else:
        raise ValueError('the weights are different. The new model is likely different from the original one.'			)

    # The model has been saved with torch.save(model) and this is bound to the exact
    # directory structure. We save the state_dict instead.
    logging.info('saving the model\'s state dictionary'			)
    torch.save(
        new_model.state_dict()					,		'./bertabs-finetuned-cnndm-extractive-abstractive-summarization/pytorch_model.bin'			)


if __name__ == "__main__":
     _UpperCamelCase						:  List[str]							       =				argparse.ArgumentParser()
     parser.add_argument(
         "--bertabs_checkpoint_path",
         default=None,
         type=str,
         required=True,
         help="Path the official PyTorch dump.",
     )
     parser.add_argument(
         "--pytorch_dump_folder_path",
         default=None,
         type=str,
         required=True,
         help="Path to the output PyTorch model.",
     )
     _UpperCamelCase						:  Optional[int]							       =				parser.parse_args()

     convert_bertabs_checkpoints(
         args.bertabs_checkpoint_path,
         args.pytorch_dump_folder_path,
     )





