





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      lowerCamelCase(						__snake_case			, unittest.TestCase	):
     '''simple docstring'''


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




     @property
     def        lowerCAmelCase__		(    self							):
         return 32




     @property
     def        lowerCAmelCase__		(    self							):
         return 32




     @property
     def        lowerCAmelCase__		(    self							):
         return self.time_input_dim * 4




     @property
     def        lowerCAmelCase__		(    self							):
         return 8




     @property
     def        lowerCAmelCase__		(    self							):
         torch.manual_seed(0							)
         _A     =							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				,     )

         _A     =							CLIPVisionModel(snake_case_							)
         return model




     @property
     def        lowerCAmelCase__		(    self							):
         _A     =							CLIPImageProcessor(
             crop_size=224				,     do_center_crop=snake_case_				,     do_normalize=snake_case_				,     do_resize=snake_case_				,     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        lowerCAmelCase__		(    self							):
         torch.manual_seed(0							)

         _A     =							{
             '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,
         }

         _A     =							PriorTransformer(**snake_case_							)
         return model




     @property
     def        lowerCAmelCase__		(    self							):
         torch.manual_seed(0							)

         _A     =							{
             '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,
             ),
         }
         _A     =							ShapERenderer(**snake_case_							)
         return model




     def        lowerCAmelCase__		(    self							):
         _A     =							self.dummy_prior
         _A     =							self.dummy_image_encoder
         _A     =							self.dummy_image_processor
         _A     =							self.dummy_renderer

         _A     =							HeunDiscreteScheduler(
             beta_schedule='exp'				,     num_train_timesteps=1024				,     prediction_type='sample'				,     use_karras_sigmas=snake_case_				,     clip_sample=snake_case_				,     clip_sample_range=1.0				,     )
         _A     =							{
             'prior': prior,
             'image_encoder': image_encoder,
             'image_processor': image_processor,
             'renderer': renderer,
             'scheduler': scheduler,
         }

         return components




     def        lowerCAmelCase__		(    self				,     snake_case_				,     snake_case_=0							):
         _A     =							floats_tensor((1, 3, 64, 64)				,     rng=random.Random(snake_case_							)							).to(snake_case_							)

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




     def        lowerCAmelCase__		(    self							):
         _A     =							'cpu'

         _A     =							self.get_dummy_components()

         _A     =							self.pipeline_class(**snake_case_							)
         _A     =							pipe.to(snake_case_							)

         pipe.set_progress_bar_config(disable=snake_case_							)

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

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

         _A     =							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        lowerCAmelCase__		(    self							):
         # NOTE: Larger batch sizes cause this test to timeout, only test on smaller batches
         self._test_inference_batch_consistent(batch_sizes=[1, 2]							)




     def        lowerCAmelCase__		(    self							):
         _A     =							torch_device == 'cpu'
         _A     =							True
         self._test_inference_batch_single_identical(
             batch_size=2				,     test_max_difference=snake_case_				,     relax_max_difference=snake_case_				,     )




     def        lowerCAmelCase__		(    self							):
         _A     =							self.get_dummy_components()
         _A     =							self.pipeline_class(**snake_case_							)
         _A     =							pipe.to(snake_case_							)
         pipe.set_progress_bar_config(disable=snake_case_							)

         _A     =							1
         _A     =							2

         _A     =							self.get_dummy_inputs(snake_case_							)

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

         _A     =							pipe(**snake_case_				,     num_images_per_prompt=snake_case_							)[0]

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







@slow
@require_torch_gpu
class      lowerCamelCase(						unittest.TestCase	):
     '''simple docstring'''


     def        lowerCAmelCase__		(    self							):
         # clean up the VRAM after each test
         super().tearDown()
         gc.collect()
         torch.cuda.empty_cache()




     def        lowerCAmelCase__		(    self							):
         _A     =							load_image(
             'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main' '/shap_e/corgi.png'							)
         _A     =							load_numpy(
             'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
             '/shap_e/test_shap_e_img2img_out.npy'							)
         _A     =							ShapEImgaImgPipeline.from_pretrained('openai/shap-e-img2img'							)
         _A     =							pipe.to(snake_case_							)
         pipe.set_progress_bar_config(disable=snake_case_							)

         _A     =							torch.Generator(device=snake_case_							).manual_seed(0							)

         _A     =							pipe(
             snake_case_				,     generator=snake_case_				,     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(snake_case_				,     snake_case_							)












from ... import PretrainedConfig


__A							:			Optional[Any]         =			{
    "sijunhe/nezha-cn-base": "https://huggingface.co/sijunhe/nezha-cn-base/resolve/main/config.json",
}







class      lowerCamelCase(						__snake_case	):
     '''simple docstring'''


     __magic_name__   =			NEZHA_PRETRAINED_CONFIG_ARCHIVE_MAP
     __magic_name__   =			'nezha'




     def __init__(    self				,     snake_case_=2_1128				,     snake_case_=768				,     snake_case_=12				,     snake_case_=12				,     snake_case_=3072				,     snake_case_="gelu"				,     snake_case_=0.1				,     snake_case_=0.1				,     snake_case_=512				,     snake_case_=64				,     snake_case_=2				,     snake_case_=0.02				,     snake_case_=1E-12				,     snake_case_=0.1				,     snake_case_=0				,     snake_case_=2				,     snake_case_=3				,     snake_case_=True				,     **snake_case_				,     ):
         super().__init__(pad_token_id=snake_case_				,     bos_token_id=snake_case_				,     eos_token_id=snake_case_				,     **snake_case_							)

         _A     =							vocab_size
         _A     =							hidden_size
         _A     =							num_hidden_layers
         _A     =							num_attention_heads
         _A     =							hidden_act
         _A     =							intermediate_size
         _A     =							hidden_dropout_prob
         _A     =							attention_probs_dropout_prob
         _A     =							max_position_embeddings
         _A     =							max_relative_position
         _A     =							type_vocab_size
         _A     =							initializer_range
         _A     =							layer_norm_eps
         _A     =							classifier_dropout
         _A     =							use_cache





