






'''simple docstring'''

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 							__magic_name__     (	UpperCAmelCase_,						unittest.TestCase				):







	"""simple docstring"""



	SCREAMING_SNAKE_CASE_					:      Dict        =  ShapEPipeline
	SCREAMING_SNAKE_CASE_					:      List[Any]        =  ['prompt']
	SCREAMING_SNAKE_CASE_					:      Tuple        =  ['prompt']
	SCREAMING_SNAKE_CASE_					:      Union[str, Any]        =  [
	    'num_images_per_prompt',
	    'num_inference_steps',
	    'generator',
	    'latents',
	    'guidance_scale',
	    'frame_size',
	    'output_type',
	    'return_dict',
	]
	SCREAMING_SNAKE_CASE_					:      int        =  False
	@property
	def        lowerCamelCase__		(							self       )						->  Dict:
							return 32
	@property
	def        lowerCamelCase__		(							self       )						->  int:
							return 32
	@property
	def        lowerCamelCase__		(							self       )						->  List[str]:
							return self.time_input_dim * 4
	@property
	def        lowerCamelCase__		(							self       )						->  str:
							return 8
	@property
	def        lowerCamelCase__		(							self       )						->  Tuple:
							lowercase_      :   str        =			CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'       )
							return tokenizer
	@property
	def        lowerCamelCase__		(							self       )						->  Optional[Any]:
							torch.manual_seed(0       )
							lowercase_      :   Optional[int]        =			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-0_5  ,      num_attention_heads=4  ,      num_hidden_layers=5  ,      pad_token_id=1  ,      vocab_size=1000  ,      )
							return CLIPTextModelWithProjection(_lowercase       )
	@property
	def        lowerCamelCase__		(							self       )						->  Any:
							torch.manual_seed(0       )

							lowercase_      :   Union[str, 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_      :   Union[str, Any]        =			PriorTransformer(**_lowercase       )
							return model
	@property
	def        lowerCamelCase__		(							self       )						->  Union[str, Any]:
							torch.manual_seed(0       )

							lowercase_      :   List[str]        =			{
							    '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_      :   Dict        =			ShapERenderer(**_lowercase       )
							return model
	def        lowerCamelCase__		(							self       )						->  Tuple:
							lowercase_      :   Optional[int]        =			self.dummy_prior
							lowercase_      :   Optional[Any]        =			self.dummy_text_encoder
							lowercase_      :   Optional[int]        =			self.dummy_tokenizer
							lowercase_      :   List[Any]        =			self.dummy_renderer

							lowercase_      :   Any        =			HeunDiscreteScheduler(
							    beta_schedule='exp'  ,      num_train_timesteps=1024  ,      prediction_type='sample'  ,      use_karras_sigmas=_lowercase  ,      clip_sample=_lowercase  ,      clip_sample_range=1.0  ,      )
							lowercase_      :   str        =			{
							    'prior': prior,
							    'text_encoder': text_encoder,
							    'tokenizer': tokenizer,
							    'renderer': renderer,
							    'scheduler': scheduler,
							}

							return components
	def        lowerCamelCase__		(							self  ,      _lowercase  ,      _lowercase=0       )						->  Tuple:
							if str(_lowercase       ).startswith('mps'       ):
													lowercase_      :   Union[str, Any]        =			torch.manual_seed(_lowercase       )
							else:
													lowercase_      :   int        =			torch.Generator(device=_lowercase       ).manual_seed(_lowercase       )
							lowercase_      :   Union[str, Any]        =			{
							    'prompt': 'horse',
							    'generator': generator,
							    'num_inference_steps': 1,
							    'frame_size': 32,
							    'output_type': 'np',
							}
							return inputs
	def        lowerCamelCase__		(							self       )						->  Dict:
							lowercase_      :   Optional[Any]        =			'cpu'

							lowercase_      :   List[Any]        =			self.get_dummy_components()

							lowercase_      :   Optional[Any]        =			self.pipeline_class(**_lowercase       )
							lowercase_      :   List[str]        =			pipe.to(_lowercase       )

							pipe.set_progress_bar_config(disable=_lowercase       )

							lowercase_      :   Dict        =			pipe(**self.get_dummy_inputs(_lowercase       )       )
							lowercase_      :   Dict        =			output.images[0]
							lowercase_      :   Optional[int]        =			image[0, -3:, -3:, -1]

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

							lowercase_      :   int        =			np.array(
							    [
							        0.00_03_92_16,
							        0.00_03_92_16,
							        0.00_03_92_16,
							        0.00_03_92_16,
							        0.00_03_92_16,
							        0.00_03_92_16,
							        0.00_03_92_16,
							        0.00_03_92_16,
							        0.00_03_92_16,
							    ]       )

							assert np.abs(image_slice.flatten() - expected_slice       ).max() < 1E-2
	def        lowerCamelCase__		(							self       )						->  List[str]:
							# 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       )						->  Union[str, Any]:
							lowercase_      :   str        =			torch_device == 'cpu'
							lowercase_      :   str        =			True

							self._test_inference_batch_single_identical(
							    batch_size=2  ,      test_max_difference=_lowercase  ,      relax_max_difference=_lowercase  ,      )







	def        lowerCamelCase__		(							self       )						->  Optional[Any]:
							lowercase_      :   int        =			self.get_dummy_components()
							lowercase_      :   Any        =			self.pipeline_class(**_lowercase       )
							lowercase_      :   Optional[int]        =			pipe.to(_lowercase       )
							pipe.set_progress_bar_config(disable=_lowercase       )

							lowercase_      :   List[Any]        =			1
							lowercase_      :   int        =			2

							lowercase_      :   Union[str, Any]        =			self.get_dummy_inputs(_lowercase       )

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

							lowercase_      :   Tuple        =			pipe(**_lowercase  ,      num_images_per_prompt=_lowercase       )[0]

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






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







	"""simple docstring"""



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







	def        lowerCamelCase__		(							self       )						->  Union[str, Any]:
							lowercase_      :   str        =			load_numpy(
							    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
							    '/shap_e/test_shap_e_np_out.npy'       )
							lowercase_      :   int        =			ShapEPipeline.from_pretrained('openai/shap-e'       )
							lowercase_      :   str        =			pipe.to(_lowercase       )
							pipe.set_progress_bar_config(disable=_lowercase       )

							lowercase_      :   Optional[Any]        =			torch.Generator(device=_lowercase       ).manual_seed(0       )

							lowercase_      :   Union[str, Any]        =			pipe(
							    'a shark'  ,      generator=_lowercase  ,      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(_lowercase  ,      _lowercase       )








'''simple docstring'''

from typing import TYPE_CHECKING

from ..utils import _LazyModule


A:      int  				=				{
    "config": [
        "EXTERNAL_DATA_FORMAT_SIZE_LIMIT",
        "OnnxConfig",
        "OnnxConfigWithPast",
        "OnnxSeq2SeqConfigWithPast",
        "PatchingSpec",
    ],
    "convert": ["export", "validate_model_outputs"],
    "features": ["FeaturesManager"],
    "utils": ["ParameterFormat", "compute_serialized_parameters_size"],
}


if TYPE_CHECKING:
						from .config import (
						    EXTERNAL_DATA_FORMAT_SIZE_LIMIT,
						    OnnxConfig,
						    OnnxConfigWithPast,
						    OnnxSeqaSeqConfigWithPast,
						    PatchingSpec,
						)
						from .convert import export, validate_model_outputs
						from .features import FeaturesManager
						from .utils import ParameterFormat, compute_serialized_parameters_size

else:
						import sys

						A:      Dict  				=				_LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)
