





'''simple docstring'''

import gc
import random
import unittest

import numpy as np
import torch
from transformers import CLIPImageProcessor, CLIPVisionConfig, CLIPVisionModel

from diffusers import HeunDiscreteScheduler, PriorTransformer, ShapEImgaImgPipeline
from diffusers.pipelines.shap_e import ShapERenderer
from diffusers.utils import floats_tensor, load_image, load_numpy, slow
from diffusers.utils.testing_utils import require_torch_gpu, torch_device

from ..test_pipelines_common import PipelineTesterMixin, assert_mean_pixel_difference


class 					a_     (				snake_case_      ,    unittest.TestCase       ):

    '''simple docstring'''
    UpperCamelCase     = ShapEImgaImgPipeline
    UpperCamelCase     = ['''image''']
    UpperCamelCase     = ['''image''']
    UpperCamelCase     = [
        '''num_images_per_prompt''',
        '''num_inference_steps''',
        '''generator''',
        '''latents''',
        '''guidance_scale''',
        '''frame_size''',
        '''output_type''',
        '''return_dict''',
    ]
    UpperCamelCase     = False


    @property
    def 			snake_case_(		self       )		->   Optional[Any]:
        return 32


    @property
    def 			snake_case_(		self       )		->   Union[str, Any]:
        return 32


    @property
    def 			snake_case_(		self       )		->   int:
        return self.time_input_dim * 4


    @property
    def 			snake_case_(		self       )		->   Tuple:
        return 8


    @property
    def 			snake_case_(		self       )		->   List[Any]:
        torch.manual_seed(0       )
        _SCREAMING_SNAKE_CASE									=			CLIPVisionConfig(
            hidden_size=self.text_embedder_hidden_size ,		image_size=64 ,		projection_dim=self.text_embedder_hidden_size ,		intermediate_size=37 ,		num_attention_heads=4 ,		num_channels=3 ,		num_hidden_layers=5 ,		patch_size=1 ,		)

        _SCREAMING_SNAKE_CASE									=			CLIPVisionModel(A       )
        return model


    @property
    def 			snake_case_(		self       )		->   str:
        _SCREAMING_SNAKE_CASE									=			CLIPImageProcessor(
            crop_size=224 ,		do_center_crop=A ,		do_normalize=A ,		do_resize=A ,		image_mean=[0.4814_5466, 0.457_8275, 0.4082_1073] ,		image_std=[0.2686_2954, 0.2613_0258, 0.2757_7711] ,		resample=3 ,		size=224 ,		)

        return image_processor


    @property
    def 			snake_case_(		self       )		->   Optional[int]:
        torch.manual_seed(0       )

        _SCREAMING_SNAKE_CASE									=			{
            """num_attention_heads""": 2,
            """attention_head_dim""": 16,
            """embedding_dim""": self.time_input_dim,
            """num_embeddings""": 32,
            """embedding_proj_dim""": self.text_embedder_hidden_size,
            """time_embed_dim""": self.time_embed_dim,
            """num_layers""": 1,
            """clip_embed_dim""": self.time_input_dim * 2,
            """additional_embeddings""": 0,
            """time_embed_act_fn""": """gelu""",
            """norm_in_type""": """layer""",
            """embedding_proj_norm_type""": """layer""",
            """encoder_hid_proj_type""": None,
            """added_emb_type""": None,
        }

        _SCREAMING_SNAKE_CASE									=			PriorTransformer(**A       )
        return model


    @property
    def 			snake_case_(		self       )		->   Any:
        torch.manual_seed(0       )

        _SCREAMING_SNAKE_CASE									=			{
            """param_shapes""": (
                (self.renderer_dim, 93),
                (self.renderer_dim, 8),
                (self.renderer_dim, 8),
                (self.renderer_dim, 8),
            ),
            """d_latent""": self.time_input_dim,
            """d_hidden""": self.renderer_dim,
            """n_output""": 12,
            """background""": (
                0.1,
                0.1,
                0.1,
            ),
        }
        _SCREAMING_SNAKE_CASE									=			ShapERenderer(**A       )
        return model


    def 			snake_case_(		self       )		->   Any:
        _SCREAMING_SNAKE_CASE									=			self.dummy_prior
        _SCREAMING_SNAKE_CASE									=			self.dummy_image_encoder
        _SCREAMING_SNAKE_CASE									=			self.dummy_image_processor
        _SCREAMING_SNAKE_CASE									=			self.dummy_renderer

        _SCREAMING_SNAKE_CASE									=			HeunDiscreteScheduler(
            beta_schedule="""exp""" ,		num_train_timesteps=1024 ,		prediction_type="""sample""" ,		use_karras_sigmas=A ,		clip_sample=A ,		clip_sample_range=1.0 ,		)
        _SCREAMING_SNAKE_CASE									=			{
            """prior""": prior,
            """image_encoder""": image_encoder,
            """image_processor""": image_processor,
            """renderer""": renderer,
            """scheduler""": scheduler,
        }

        return components


    def 			snake_case_(		self ,		A ,		A=0       )		->   Tuple:
        _SCREAMING_SNAKE_CASE									=			floats_tensor((1, 3, 64, 64) ,		rng=random.Random(A       )       ).to(A       )

        if str(A       ).startswith("""mps"""       ):
            _SCREAMING_SNAKE_CASE									=			torch.manual_seed(A       )
        else:
            _SCREAMING_SNAKE_CASE									=			torch.Generator(device=A       ).manual_seed(A       )
        _SCREAMING_SNAKE_CASE									=			{
            """image""": input_image,
            """generator""": generator,
            """num_inference_steps""": 1,
            """frame_size""": 32,
            """output_type""": """np""",
        }
        return inputs


    def 			snake_case_(		self       )		->   Any:
        _SCREAMING_SNAKE_CASE									=			"""cpu"""

        _SCREAMING_SNAKE_CASE									=			self.get_dummy_components()

        _SCREAMING_SNAKE_CASE									=			self.pipeline_class(**A       )
        _SCREAMING_SNAKE_CASE									=			pipe.to(A       )

        pipe.set_progress_bar_config(disable=A       )

        _SCREAMING_SNAKE_CASE									=			pipe(**self.get_dummy_inputs(A       )       )
        _SCREAMING_SNAKE_CASE									=			output.images[0]
        _SCREAMING_SNAKE_CASE									=			image[0, -3:, -3:, -1]

        assert image.shape == (20, 32, 32, 3)

        _SCREAMING_SNAKE_CASE									=			np.array(
            [
                0.0003_9216,
                0.0003_9216,
                0.0003_9216,
                0.0003_9216,
                0.0003_9216,
                0.0003_9216,
                0.0003_9216,
                0.0003_9216,
                0.0003_9216,
            ]       )

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


    def 			snake_case_(		self       )		->   List[Any]:
        # NOTE: Larger batch sizes cause this test to timeout, only test on smaller batches
        self._test_inference_batch_consistent(batch_sizes=[1, 2]       )


    def 			snake_case_(		self       )		->   Optional[Any]:
        _SCREAMING_SNAKE_CASE									=			torch_device == """cpu"""
        _SCREAMING_SNAKE_CASE									=			True
        self._test_inference_batch_single_identical(
            batch_size=2 ,		test_max_difference=A ,		relax_max_difference=A ,		)





    def 			snake_case_(		self       )		->   Optional[int]:
        _SCREAMING_SNAKE_CASE									=			self.get_dummy_components()
        _SCREAMING_SNAKE_CASE									=			self.pipeline_class(**A       )
        _SCREAMING_SNAKE_CASE									=			pipe.to(A       )
        pipe.set_progress_bar_config(disable=A       )

        _SCREAMING_SNAKE_CASE									=			1
        _SCREAMING_SNAKE_CASE									=			2

        _SCREAMING_SNAKE_CASE									=			self.get_dummy_inputs(A       )

        for key in inputs.keys():
            if key in self.batch_params:
                _SCREAMING_SNAKE_CASE									=			batch_size * [inputs[key]]

        _SCREAMING_SNAKE_CASE									=			pipe(**A ,		num_images_per_prompt=A       )[0]

        assert images.shape[0] == batch_size * num_images_per_prompt




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

    '''simple docstring'''


    def 			snake_case_(		self       )		->   Union[str, Any]:
        # clean up the VRAM after each test
        super().tearDown()
        gc.collect()
        torch.cuda.empty_cache()





    def 			snake_case_(		self       )		->   List[Any]:
        _SCREAMING_SNAKE_CASE									=			load_image(
            """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main""" """/shap_e/corgi.png"""       )
        _SCREAMING_SNAKE_CASE									=			load_numpy(
            """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"""
            """/shap_e/test_shap_e_img2img_out.npy"""       )
        _SCREAMING_SNAKE_CASE									=			ShapEImgaImgPipeline.from_pretrained("""openai/shap-e-img2img"""       )
        _SCREAMING_SNAKE_CASE									=			pipe.to(A       )
        pipe.set_progress_bar_config(disable=A       )

        _SCREAMING_SNAKE_CASE									=			torch.Generator(device=A       ).manual_seed(0       )

        _SCREAMING_SNAKE_CASE									=			pipe(
            A ,		generator=A ,		guidance_scale=3.0 ,		num_inference_steps=64 ,		frame_size=64 ,		output_type="""np""" ,		).images[0]

        assert images.shape == (20, 64, 64, 3)

        assert_mean_pixel_difference(A ,		A       )










'''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)
lowercase_					= logging.getLogger(__name__)


lowercase_					= """Hello world! cécé herlolip"""


lowercase_					= 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 lowerCamelCase						(     __lowerCamelCase :    Optional[Any]					,       __lowerCamelCase :    Union[str, Any]    )						->List[Any]:
    _SCREAMING_SNAKE_CASE									=			BertAbsConfig(
        temp_dir="""."""					,       finetune_bert=__lowerCamelCase					,       large=__lowerCamelCase					,       share_emb=__lowerCamelCase					,       use_bert_emb=__lowerCamelCase					,       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					,       )
    _SCREAMING_SNAKE_CASE									=			torch.load(__lowerCamelCase					,       lambda __lowerCamelCase					,       __lowerCamelCase    : storage    )
    _SCREAMING_SNAKE_CASE									=			AbsSummarizer(__lowerCamelCase					,       torch.device("""cpu"""    )					,       __lowerCamelCase    )
    original.eval()

    _SCREAMING_SNAKE_CASE									=			BertAbsSummarizer(__lowerCamelCase					,       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"""    )
    _SCREAMING_SNAKE_CASE									=			BertTokenizer.from_pretrained("""bert-base-uncased"""    )

    # prepare the model inputs
    _SCREAMING_SNAKE_CASE									=			tokenizer.encode("""This is sample éàalj'-."""    )
    encoder_input_ids.extend([tokenizer.pad_token_id] * (512 - len(__lowerCamelCase    ))    )
    _SCREAMING_SNAKE_CASE									=			torch.tensor(__lowerCamelCase    ).unsqueeze(0    )
    _SCREAMING_SNAKE_CASE									=			tokenizer.encode("""This is sample 3 éàalj'-."""    )
    decoder_input_ids.extend([tokenizer.pad_token_id] * (512 - len(__lowerCamelCase    ))    )
    _SCREAMING_SNAKE_CASE									=			torch.tensor(__lowerCamelCase    ).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
    _SCREAMING_SNAKE_CASE									=			encoder_input_ids
    _SCREAMING_SNAKE_CASE									=			decoder_input_ids
    _SCREAMING_SNAKE_CASE									=			_SCREAMING_SNAKE_CASE									=			None
    _SCREAMING_SNAKE_CASE									=			None
    _SCREAMING_SNAKE_CASE									=			_SCREAMING_SNAKE_CASE									=			None
    _SCREAMING_SNAKE_CASE									=			_SCREAMING_SNAKE_CASE									=			None
    _SCREAMING_SNAKE_CASE									=			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
    _SCREAMING_SNAKE_CASE									=			original(__lowerCamelCase					,       __lowerCamelCase					,       __lowerCamelCase					,       __lowerCamelCase					,       __lowerCamelCase					,       __lowerCamelCase					,       __lowerCamelCase    )[0]
    _SCREAMING_SNAKE_CASE									=			original.generator(__lowerCamelCase    )

    _SCREAMING_SNAKE_CASE									=			new_model(
        __lowerCamelCase					,       __lowerCamelCase					,       __lowerCamelCase					,       __lowerCamelCase					,       __lowerCamelCase    )[0]
    _SCREAMING_SNAKE_CASE									=			new_model.generator(__lowerCamelCase    )

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

    _SCREAMING_SNAKE_CASE									=			torch.allclose(__lowerCamelCase					,       __lowerCamelCase					,       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__":
      lowercase_					= 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.""",
      )
      lowercase_					= parser.parse_args()

      convert_bertabs_checkpoints(
          args.bertabs_checkpoint_path,
          args.pytorch_dump_folder_path,
      )



