


import gc
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModelWithProjection, CLIPTokenizer

from diffusers import HeunDiscreteScheduler, PriorTransformer, ShapEPipeline
from diffusers.pipelines.shap_e import ShapERenderer
from diffusers.utils import 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				):
 __lowercase			:							List[str]       					=			ShapEPipeline
 __lowercase			:							Tuple       					=			['prompt']
 __lowercase			:							Optional[int]       					=			['prompt']
 __lowercase			:							Any       					=			[
     'num_images_per_prompt',
     'num_inference_steps',
     'generator',
     'latents',
     'guidance_scale',
     'frame_size',
     'output_type',
     'return_dict',
 ]
 __lowercase			:							List[Any]       					=			False
 @property
 def 		SCREAMING_SNAKE_CASE__       (       self			) ->    Optional[Any]:




        '''simple docstring'''
        return 32
 @property
 def 		SCREAMING_SNAKE_CASE__       (       self			) ->    List[Any]:




        '''simple docstring'''
        return 32
 @property
 def 		SCREAMING_SNAKE_CASE__       (       self			) ->    str:




        '''simple docstring'''
        return self.time_input_dim * 4
 @property
 def 		SCREAMING_SNAKE_CASE__       (       self			) ->    str:




        '''simple docstring'''
        return 8
 @property
 def 		SCREAMING_SNAKE_CASE__       (       self			) ->    Optional[int]:




        '''simple docstring'''
        lowercase__:       List[str]        =       CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'			)
        return tokenizer
 @property
 def 		SCREAMING_SNAKE_CASE__       (       self			) ->    Optional[Any]:




        '''simple docstring'''
        torch.manual_seed(0			)
        lowercase__:       Union[str, Any]        =       CLIPTextConfig(
            bos_token_id=0    ,			eos_token_id=2    ,			hidden_size=self.text_embedder_hidden_size    ,			projection_dim=self.text_embedder_hidden_size    ,			intermediate_size=37    ,			layer_norm_eps=1E-05    ,			num_attention_heads=4    ,			num_hidden_layers=5    ,			pad_token_id=1    ,			vocab_size=1_000    ,			)
        return CLIPTextModelWithProjection(lowerCAmelCase__			)
 @property
 def 		SCREAMING_SNAKE_CASE__       (       self			) ->    Optional[int]:




        '''simple docstring'''
        torch.manual_seed(0			)

        lowercase__:       List[Any]        =       {
            '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',
            'encoder_hid_proj_type': None,
            'added_emb_type': None,
        }

        lowercase__:       List[str]        =       PriorTransformer(**lowerCAmelCase__			)
        return model
 @property
 def 		SCREAMING_SNAKE_CASE__       (       self			) ->    Optional[int]:




        '''simple docstring'''
        torch.manual_seed(0			)

        lowercase__:       Any        =       {
            '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,
            ),
        }
        lowercase__:       Any        =       ShapERenderer(**lowerCAmelCase__			)
        return model
 def 		SCREAMING_SNAKE_CASE__       (       self			) ->    Optional[Any]:




        '''simple docstring'''
        lowercase__:       List[str]        =       self.dummy_prior
        lowercase__:       List[Any]        =       self.dummy_text_encoder
        lowercase__:       int        =       self.dummy_tokenizer
        lowercase__:       Optional[Any]        =       self.dummy_renderer

        lowercase__:       Tuple        =       HeunDiscreteScheduler(
            beta_schedule='exp'    ,			num_train_timesteps=1_024    ,			prediction_type='sample'    ,			use_karras_sigmas=lowerCAmelCase__    ,			clip_sample=lowerCAmelCase__    ,			clip_sample_range=1.0    ,			)
        lowercase__:       int        =       {
            'prior': prior,
            'text_encoder': text_encoder,
            'tokenizer': tokenizer,
            'renderer': renderer,
            'scheduler': scheduler,
        }

        return components
 def 		SCREAMING_SNAKE_CASE__       (       self    ,			lowerCAmelCase__    ,			lowerCAmelCase__=0			) ->    Optional[int]:




        '''simple docstring'''
        if str(lowerCAmelCase__			).startswith('mps'			):
               lowercase__:       int        =       torch.manual_seed(lowerCAmelCase__			)
        else:
               lowercase__:       Dict        =       torch.Generator(device=lowerCAmelCase__			).manual_seed(lowerCAmelCase__			)
        lowercase__:       Any        =       {
            'prompt': 'horse',
            'generator': generator,
            'num_inference_steps': 1,
            'frame_size': 32,
            'output_type': 'np',
        }
        return inputs
 def 		SCREAMING_SNAKE_CASE__       (       self			) ->    Tuple:




        '''simple docstring'''
        lowercase__:       Optional[int]        =       'cpu'

        lowercase__:       Any        =       self.get_dummy_components()

        lowercase__:       str        =       self.pipeline_class(**lowerCAmelCase__			)
        lowercase__:       Dict        =       pipe.to(lowerCAmelCase__			)

        pipe.set_progress_bar_config(disable=lowerCAmelCase__			)

        lowercase__:       Optional[Any]        =       pipe(**self.get_dummy_inputs(lowerCAmelCase__			)			)
        lowercase__:       List[Any]        =       output.images[0]
        lowercase__:       Optional[int]        =       image[0, -3:, -3:, -1]

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

        lowercase__:       Any        =       np.array(
            [
                0.0_0_0_3_9_2_1_6,
                0.0_0_0_3_9_2_1_6,
                0.0_0_0_3_9_2_1_6,
                0.0_0_0_3_9_2_1_6,
                0.0_0_0_3_9_2_1_6,
                0.0_0_0_3_9_2_1_6,
                0.0_0_0_3_9_2_1_6,
                0.0_0_0_3_9_2_1_6,
                0.0_0_0_3_9_2_1_6,
            ]			)

        assert np.abs(image_slice.flatten() - expected_slice			).max() < 1E-2
 def 		SCREAMING_SNAKE_CASE__       (       self			) ->    Any:




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




        '''simple docstring'''
        lowercase__:       Tuple        =       torch_device == 'cpu'
        lowercase__:       str        =       True

        self._test_inference_batch_single_identical(
            batch_size=2    ,			test_max_difference=lowerCAmelCase__    ,			relax_max_difference=lowerCAmelCase__    ,			)
 def 		SCREAMING_SNAKE_CASE__       (       self			) ->    int:




        '''simple docstring'''
        lowercase__:       Any        =       self.get_dummy_components()
        lowercase__:       Union[str, Any]        =       self.pipeline_class(**lowerCAmelCase__			)
        lowercase__:       List[str]        =       pipe.to(lowerCAmelCase__			)
        pipe.set_progress_bar_config(disable=lowerCAmelCase__			)

        lowercase__:       Tuple        =       1
        lowercase__:       Optional[int]        =       2

        lowercase__:       Optional[Any]        =       self.get_dummy_inputs(lowerCAmelCase__			)

        for key in inputs.keys():
               if key in self.batch_params:
                      lowercase__:       Union[str, Any]        =       batch_size * [inputs[key]]

        lowercase__:       Optional[int]        =       pipe(**lowerCAmelCase__    ,			num_images_per_prompt=lowerCAmelCase__			)[0]

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





@slow
@require_torch_gpu
class   __a      (   unittest.TestCase				):
 def 		SCREAMING_SNAKE_CASE__       (       self			) ->    List[str]:




        '''simple docstring'''
        # clean up the VRAM after each test
        super().tearDown()
        gc.collect()
        torch.cuda.empty_cache()
 def 		SCREAMING_SNAKE_CASE__       (       self			) ->    Dict:




        '''simple docstring'''
        lowercase__:       Dict        =       load_numpy(
            'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
            '/shap_e/test_shap_e_np_out.npy'			)
        lowercase__:       str        =       ShapEPipeline.from_pretrained('openai/shap-e'			)
        lowercase__:       Tuple        =       pipe.to(lowerCAmelCase__			)
        pipe.set_progress_bar_config(disable=lowerCAmelCase__			)

        lowercase__:       Any        =       torch.Generator(device=lowerCAmelCase__			).manual_seed(0			)

        lowercase__:       Tuple        =       pipe(
            'a shark'    ,			generator=lowerCAmelCase__    ,			guidance_scale=15.0    ,			num_inference_steps=64    ,			frame_size=64    ,			output_type='np'    ,			).images[0]

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

        assert_mean_pixel_difference(lowerCAmelCase__    ,			lowerCAmelCase__			)





from ...configuration_utils import PretrainedConfig
from ...utils import logging


__lowerCAmelCase			=							logging.get_logger(__name__)

__lowerCAmelCase			=							{
    '''google/pegasus-large''': '''https://huggingface.co/google/pegasus-large/resolve/main/config.json''',
    # See all PEGASUS models at https://huggingface.co/models?filter=pegasus
}





class   __a      (   __UpperCamelCase				):
 __lowercase			:							Any       					=			'pegasus'
 __lowercase			:							Union[str, Any]       					=			['past_key_values']
 __lowercase			:							Any       					=			{'num_attention_heads': 'encoder_attention_heads', 'hidden_size': 'd_model'}
 def __init__(       self    ,			lowerCAmelCase__=50_265    ,			lowerCAmelCase__=1_024    ,			lowerCAmelCase__=12    ,			lowerCAmelCase__=4_096    ,			lowerCAmelCase__=16    ,			lowerCAmelCase__=12    ,			lowerCAmelCase__=4_096    ,			lowerCAmelCase__=16    ,			lowerCAmelCase__=0.0    ,			lowerCAmelCase__=0.0    ,			lowerCAmelCase__=True    ,			lowerCAmelCase__=True    ,			lowerCAmelCase__="gelu"    ,			lowerCAmelCase__=1_024    ,			lowerCAmelCase__=0.1    ,			lowerCAmelCase__=0.0    ,			lowerCAmelCase__=0.0    ,			lowerCAmelCase__=0.0_2    ,			lowerCAmelCase__=0    ,			lowerCAmelCase__=False    ,			lowerCAmelCase__=0    ,			lowerCAmelCase__=1    ,			lowerCAmelCase__=1    ,			**lowerCAmelCase__    ,			) ->    Union[str, Any]:




        '''simple docstring'''
        lowercase__:       int        =       vocab_size
        lowercase__:       Optional[int]        =       max_position_embeddings
        lowercase__:       List[str]        =       d_model
        lowercase__:       Optional[Any]        =       encoder_ffn_dim
        lowercase__:       Optional[Any]        =       encoder_layers
        lowercase__:       Union[str, Any]        =       encoder_attention_heads
        lowercase__:       Optional[int]        =       decoder_ffn_dim
        lowercase__:       Tuple        =       decoder_layers
        lowercase__:       Union[str, Any]        =       decoder_attention_heads
        lowercase__:       Dict        =       dropout
        lowercase__:       List[str]        =       attention_dropout
        lowercase__:       List[str]        =       activation_dropout
        lowercase__:       Optional[int]        =       activation_function
        lowercase__:       Dict        =       init_std
        lowercase__:       Optional[Any]        =       encoder_layerdrop
        lowercase__:       List[str]        =       decoder_layerdrop
        lowercase__:       Union[str, Any]        =       use_cache
        lowercase__:       Any        =       encoder_layers
        lowercase__:       List[str]        =       scale_embedding  # scale factor will be sqrt(d_model) if True
        super().__init__(
            pad_token_id=lowerCAmelCase__    ,			eos_token_id=lowerCAmelCase__    ,			is_encoder_decoder=lowerCAmelCase__    ,			decoder_start_token_id=lowerCAmelCase__    ,			forced_eos_token_id=lowerCAmelCase__    ,			**lowerCAmelCase__    ,			)
 @property
 def 		SCREAMING_SNAKE_CASE__       (       self			) ->    int:




        '''simple docstring'''
        return self.encoder_attention_heads
 @property
 def 		SCREAMING_SNAKE_CASE__       (       self			) ->    int:




        '''simple docstring'''
        return self.d_model


