
"""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_    (			_UpperCAmelCase  ,    unittest.TestCase						):
    a   :					Optional[int]			     =					ShapEImgaImgPipeline
    a   :					Any			     =					['image']
    a   :					str			     =					['image']
    a   :					Optional[Any]			     =					[
        'num_images_per_prompt',
        'num_inference_steps',
        'generator',
        'latents',
        'guidance_scale',
        'frame_size',
        'output_type',
        'return_dict',
    ]
    a   :					Union[str, Any]			     =					False







    @property
    def   _snake_case       (      self      :							str  )		->Any:






        '''simple docstring'''






        return 32







    @property
    def   _snake_case       (      self      :							int  )		->Optional[int]:






        '''simple docstring'''






        return 32







    @property
    def   _snake_case       (      self      :							Optional[Any]  )		->str:






        '''simple docstring'''






        return self.time_input_dim * 4







    @property
    def   _snake_case       (      self      :							Optional[int]  )		->Optional[int]:






        '''simple docstring'''






        return 8







    @property
    def   _snake_case       (      self      :							Union[str, Any]  )		->Tuple:






        '''simple docstring'''






        torch.manual_seed(0  )
        _UpperCAmelCase        = 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  ,       )

        _UpperCAmelCase        = CLIPVisionModel(__UpperCamelCase  )
        return model







    @property
    def   _snake_case       (      self      :							Optional[int]  )		->Optional[Any]:






        '''simple docstring'''






        _UpperCAmelCase        = CLIPImageProcessor(
            crop_size=2_24  ,       do_center_crop=__UpperCamelCase  ,       do_normalize=__UpperCamelCase  ,       do_resize=__UpperCamelCase  ,       image_mean=[0.4_8_1_4_5_4_6_6, 0.4_5_7_8_2_7_5, 0.4_0_8_2_1_0_7_3]  ,       image_std=[0.2_6_8_6_2_9_5_4, 0.2_6_1_3_0_2_5_8, 0.2_7_5_7_7_7_1_1]  ,       resample=3  ,       size=2_24  ,       )

        return image_processor







    @property
    def   _snake_case       (      self      :							Union[str, Any]  )		->Optional[Any]:






        '''simple docstring'''






        torch.manual_seed(0  )

        _UpperCAmelCase        = {
            """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,
        }

        _UpperCAmelCase        = PriorTransformer(**__UpperCamelCase  )
        return model







    @property
    def   _snake_case       (      self      :							Dict  )		->Union[str, Any]:






        '''simple docstring'''






        torch.manual_seed(0  )

        _UpperCAmelCase        = {
            """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,
            ),
        }
        _UpperCAmelCase        = ShapERenderer(**__UpperCamelCase  )
        return model







    def   _snake_case       (      self      :							List[str]  )		->Optional[Any]:






        '''simple docstring'''






        _UpperCAmelCase        = self.dummy_prior
        _UpperCAmelCase        = self.dummy_image_encoder
        _UpperCAmelCase        = self.dummy_image_processor
        _UpperCAmelCase        = self.dummy_renderer

        _UpperCAmelCase        = HeunDiscreteScheduler(
            beta_schedule="""exp"""  ,       num_train_timesteps=10_24  ,       prediction_type="""sample"""  ,       use_karras_sigmas=__UpperCamelCase  ,       clip_sample=__UpperCamelCase  ,       clip_sample_range=1.0  ,       )
        _UpperCAmelCase        = {
            """prior""": prior,
            """image_encoder""": image_encoder,
            """image_processor""": image_processor,
            """renderer""": renderer,
            """scheduler""": scheduler,
        }

        return components







    def   _snake_case       (      self      :							int  ,       __UpperCamelCase      :							Tuple  ,       __UpperCamelCase      :							Tuple=0  )		->Any:






        '''simple docstring'''






        _UpperCAmelCase        = floats_tensor((1, 3, 64, 64)  ,       rng=random.Random(__UpperCamelCase  )  ).to(__UpperCamelCase  )

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







    def   _snake_case       (      self      :							Union[str, Any]  )		->List[Any]:






        '''simple docstring'''






        _UpperCAmelCase        = """cpu"""

        _UpperCAmelCase        = self.get_dummy_components()

        _UpperCAmelCase        = self.pipeline_class(**__UpperCamelCase  )
        _UpperCAmelCase        = pipe.to(__UpperCamelCase  )

        pipe.set_progress_bar_config(disable=__UpperCamelCase  )

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

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

        _UpperCAmelCase        = 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   _snake_case       (      self      :							Union[str, Any]  )		->Tuple:






        '''simple docstring'''






        self._test_inference_batch_consistent(batch_sizes=[1, 2]  )







    def   _snake_case       (      self      :							int  )		->List[Any]:






        '''simple docstring'''






        _UpperCAmelCase        = torch_device == """cpu"""
        _UpperCAmelCase        = True
        self._test_inference_batch_single_identical(
            batch_size=2  ,       test_max_difference=__UpperCamelCase  ,       relax_max_difference=__UpperCamelCase  ,       )







    def   _snake_case       (      self      :							List[Any]  )		->Union[str, Any]:






        '''simple docstring'''






        _UpperCAmelCase        = self.get_dummy_components()
        _UpperCAmelCase        = self.pipeline_class(**__UpperCamelCase  )
        _UpperCAmelCase        = pipe.to(__UpperCamelCase  )
        pipe.set_progress_bar_config(disable=__UpperCamelCase  )

        _UpperCAmelCase        = 1
        _UpperCAmelCase        = 2

        _UpperCAmelCase        = self.get_dummy_inputs(__UpperCamelCase  )

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

        _UpperCAmelCase        = pipe(**__UpperCamelCase  ,       num_images_per_prompt=__UpperCamelCase  )[0]

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







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







    def   _snake_case       (      self      :							List[Any]  )		->Tuple:






        '''simple docstring'''






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







    def   _snake_case       (      self      :							str  )		->Any:






        '''simple docstring'''






        _UpperCAmelCase        = load_image(
            """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main""" """/shap_e/corgi.png"""  )
        _UpperCAmelCase        = load_numpy(
            """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"""
            """/shap_e/test_shap_e_img2img_out.npy"""  )
        _UpperCAmelCase        = ShapEImgaImgPipeline.from_pretrained("""openai/shap-e-img2img"""  )
        _UpperCAmelCase        = pipe.to(__UpperCamelCase  )
        pipe.set_progress_bar_config(disable=__UpperCamelCase  )

        _UpperCAmelCase        = torch.Generator(device=__UpperCamelCase  ).manual_seed(0  )

        _UpperCAmelCase        = pipe(
            __UpperCamelCase  ,       generator=__UpperCamelCase  ,       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(__UpperCamelCase  ,       __UpperCamelCase  )

"""simple docstring"""



from collections import OrderedDict
from typing import Mapping

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig
from ...utils import logging


a       :						List[Any]       		=		logging.get_logger(__name__)

a       :						str       		=		{
    '''camembert-base''': '''https://huggingface.co/camembert-base/resolve/main/config.json''',
    '''umberto-commoncrawl-cased-v1''': (
        '''https://huggingface.co/Musixmatch/umberto-commoncrawl-cased-v1/resolve/main/config.json'''
    ),
    '''umberto-wikipedia-uncased-v1''': (
        '''https://huggingface.co/Musixmatch/umberto-wikipedia-uncased-v1/resolve/main/config.json'''
    ),
}




class 		a_    (			_UpperCAmelCase						):
    a   :					Any			     =					'camembert'







    def __init__(      self      :							Dict  ,       __UpperCamelCase      :							Union[str, Any]=3_05_22  ,       __UpperCamelCase      :							Tuple=7_68  ,       __UpperCamelCase      :							Optional[Any]=12  ,       __UpperCamelCase      :							int=12  ,       __UpperCamelCase      :							int=30_72  ,       __UpperCamelCase      :							int="gelu"  ,       __UpperCamelCase      :							Dict=0.1  ,       __UpperCamelCase      :							Dict=0.1  ,       __UpperCamelCase      :							List[Any]=5_12  ,       __UpperCamelCase      :							Optional[int]=2  ,       __UpperCamelCase      :							Dict=0.0_2  ,       __UpperCamelCase      :							Optional[int]=1e-12  ,       __UpperCamelCase      :							Optional[Any]=1  ,       __UpperCamelCase      :							int=0  ,       __UpperCamelCase      :							Optional[int]=2  ,       __UpperCamelCase      :							Optional[int]="absolute"  ,       __UpperCamelCase      :							int=True  ,       __UpperCamelCase      :							Dict=None  ,       **__UpperCamelCase      :							List[str]  ,       )		->Union[str, Any]:






        '''simple docstring'''






        super().__init__(pad_token_id=__UpperCamelCase  ,       bos_token_id=__UpperCamelCase  ,       eos_token_id=__UpperCamelCase  ,       **__UpperCamelCase  )

        _UpperCAmelCase        = vocab_size
        _UpperCAmelCase        = hidden_size
        _UpperCAmelCase        = num_hidden_layers
        _UpperCAmelCase        = num_attention_heads
        _UpperCAmelCase        = hidden_act
        _UpperCAmelCase        = intermediate_size
        _UpperCAmelCase        = hidden_dropout_prob
        _UpperCAmelCase        = attention_probs_dropout_prob
        _UpperCAmelCase        = max_position_embeddings
        _UpperCAmelCase        = type_vocab_size
        _UpperCAmelCase        = initializer_range
        _UpperCAmelCase        = layer_norm_eps
        _UpperCAmelCase        = position_embedding_type
        _UpperCAmelCase        = use_cache
        _UpperCAmelCase        = classifier_dropout







class 		a_    (			_UpperCAmelCase						):







    @property
    def   _snake_case       (      self      :							Optional[int]  )		->Mapping[str, Mapping[int, str]]:






        '''simple docstring'''






        if self.task == "multiple-choice":
            _UpperCAmelCase        = {0: """batch""", 1: """choice""", 2: """sequence"""}
        else:
            _UpperCAmelCase        = {0: """batch""", 1: """sequence"""}
        return OrderedDict(
            [
                ("""input_ids""", dynamic_axis),
                ("""attention_mask""", dynamic_axis),
            ]  )