




'''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   __UpperCAmelCase				(     lowerCAmelCase			,unittest.TestCase				):



					'''simple docstring'''





					_UpperCamelCase							=		ShapEPipeline
					_UpperCamelCase							=		["""prompt"""]
					_UpperCamelCase							=		["""prompt"""]
					_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     :						List[Any])    ->							str:
								return 32



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



					@property
					def     __snake_case						(	self     :						Optional[int])    ->							List[Any]:
								return self.time_input_dim * 4



					@property
					def     __snake_case						(	self     :						Any)    ->							Any:
								return 8



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



					@property
					def     __snake_case						(	self     :						str)    ->							Optional[int]:
								torch.manual_seed(0)
								A_									= 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(_lowercase)



					@property
					def     __snake_case						(	self     :						Optional[Any])    ->							Optional[int]:
								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',
								    'encoder_hid_proj_type': None,
								    'added_emb_type': None,
								}

								A_									= PriorTransformer(**_lowercase)
								return model



					@property
					def     __snake_case						(	self     :						Any)    ->							Optional[int]:
								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(**_lowercase)
								return model



					def     __snake_case						(	self     :						Optional[int])    ->							Optional[Any]:
								A_									= self.dummy_prior
								A_									= self.dummy_text_encoder
								A_									= self.dummy_tokenizer
								A_									= self.dummy_renderer

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

								return components



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



					def     __snake_case						(	self     :						List[Any])    ->							Optional[int]:
								A_									= 'cpu'

								A_									= self.get_dummy_components()

								A_									= self.pipeline_class(**_lowercase)
								A_									= pipe.to(_lowercase)

								pipe.set_progress_bar_config(disable=_lowercase)

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

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

								A_									= 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     __snake_case						(	self     :						Tuple)    ->							Optional[int]:
								# 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     :						Tuple)    ->							List[Any]:
								A_									= torch_device == 'cpu'
								A_									= True

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



					def     __snake_case						(	self     :						int)    ->							int:
								A_									= self.get_dummy_components()
								A_									= self.pipeline_class(**_lowercase)
								A_									= pipe.to(_lowercase)
								pipe.set_progress_bar_config(disable=_lowercase)

								A_									= 1
								A_									= 2

								A_									= self.get_dummy_inputs(_lowercase)

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

								A_									= pipe(**_lowercase      ,       num_images_per_prompt=_lowercase)[0]

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





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



					'''simple docstring'''





					def     __snake_case						(	self     :						int)    ->							List[Any]:
								# clean up the VRAM after each test
								super().tearDown()
								gc.collect()
								torch.cuda.empty_cache()



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

								A_									= torch.Generator(device=_lowercase).manual_seed(0)

								A_									= 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'''
import pytest

from datasets.splits import SplitDict, SplitInfo
from datasets.utils.py_utils import asdict





@pytest.mark.parametrize(
    'split_dict'       ,[
        SplitDict(),
        SplitDict({'train': SplitInfo(name='train'       ,num_bytes=1337       ,num_examples=42       ,dataset_name='my_dataset'     )}     ),
        SplitDict({'train': SplitInfo(name='train'       ,num_bytes=1337       ,num_examples=42     )}     ),
        SplitDict({'train': SplitInfo()}     ),
    ]       ,)
def 			lowerCamelCase(							SCREAMING_SNAKE_CASE_     )		->		List[str]:
			A_									= split_dict._to_yaml_list()
			assert len(SCREAMING_SNAKE_CASE_     ) == len(SCREAMING_SNAKE_CASE_     )
			A_									= SplitDict._from_yaml_list(SCREAMING_SNAKE_CASE_     )
			for split_name, split_info in split_dict.items():
						# dataset_name field is deprecated, and is therefore not part of the YAML dump
						A_									= None
						# the split name of split_dict takes over the name of the split info object
						A_									= split_name
			assert split_dict == reloaded






@pytest.mark.parametrize(
    'split_info'       ,[SplitInfo(), SplitInfo(dataset_name=SCREAMING_SNAKE_CASE_     ), SplitInfo(dataset_name='my_dataset'     )]     )
def 			lowerCamelCase(							SCREAMING_SNAKE_CASE_     )		->		Optional[Any]:
			# For backward compatibility, we need asdict(split_dict) to return split info dictrionaries with the "dataset_name"
			# field even if it's deprecated. This way old versionso of `datasets` can still reload dataset_infos.json files
			A_									= asdict(SplitDict({'train': split_info}     )     )
			assert "dataset_name" in split_dict_asdict["train"]
			assert split_dict_asdict["train"]["dataset_name"] == split_info.dataset_name




