






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







	"""simple docstring"""



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

							lowercase_      :   Tuple        =			{
							    "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_      :   Dict        =			PriorTransformer(**snake_case_       )
							return model
	@property
	def        lowerCamelCase__		(							self       )						->  Optional[Any]:
							torch.manual_seed(0       )

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

							lowercase_      :   Dict        =			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  ,      )
							lowercase_      :   List[str]        =			{
							    "prior": prior,
							    "text_encoder": text_encoder,
							    "tokenizer": tokenizer,
							    "renderer": renderer,
							    "scheduler": scheduler,
							}

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

							lowercase_      :   int        =			self.get_dummy_components()

							lowercase_      :   Tuple        =			self.pipeline_class(**snake_case_       )
							lowercase_      :   Any        =			pipe.to(snake_case_       )

							pipe.set_progress_bar_config(disable=snake_case_       )

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

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

							lowercase_      :   List[str]        =			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[Any]:
							# 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       )						->  Any:
							lowercase_      :   List[Any]        =			torch_device == "cpu"
							lowercase_      :   Tuple        =			True

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







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

							lowercase_      :   int        =			1
							lowercase_      :   Dict        =			2

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

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

							lowercase_      :   Optional[Any]        =			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 							__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       )						->  Any:
							lowercase_      :   List[Any]        =			load_numpy(
							    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
							    '/shap_e/test_shap_e_np_out.npy'       )
							lowercase_      :   Any        =			ShapEPipeline.from_pretrained('openai/shap-e'       )
							lowercase_      :   List[Any]        =			pipe.to(snake_case_       )
							pipe.set_progress_bar_config(disable=snake_case_       )

							lowercase_      :   str        =			torch.Generator(device=snake_case_       ).manual_seed(0       )

							lowercase_      :   Dict        =			pipe(
							    'a shark'  ,      generator=snake_case_  ,      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(snake_case_  ,      snake_case_       )








'''simple docstring'''

import unittest

from diffusers.pipelines.pipeline_utils import is_safetensors_compatible






class 							__magic_name__     (	unittest.TestCase				):







	"""simple docstring"""



	def        lowerCamelCase__		(							self       )						->  Optional[Any]:
							lowercase_      :   Tuple        =			[
							    'safety_checker/pytorch_model.bin',
							    'safety_checker/model.safetensors',
							    'vae/diffusion_pytorch_model.bin',
							    'vae/diffusion_pytorch_model.safetensors',
							    'text_encoder/pytorch_model.bin',
							    'text_encoder/model.safetensors',
							    'unet/diffusion_pytorch_model.bin',
							    'unet/diffusion_pytorch_model.safetensors',
							]
							self.assertTrue(is_safetensors_compatible(_lowercase       )       )
	def        lowerCamelCase__		(							self       )						->  Dict:
							lowercase_      :   Tuple        =			[
							    'unet/diffusion_pytorch_model.bin',
							    'unet/diffusion_pytorch_model.safetensors',
							]
							self.assertTrue(is_safetensors_compatible(_lowercase       )       )
	def        lowerCamelCase__		(							self       )						->  str:
							lowercase_      :   int        =			[
							    'safety_checker/pytorch_model.bin',
							    'safety_checker/model.safetensors',
							    'vae/diffusion_pytorch_model.bin',
							    'vae/diffusion_pytorch_model.safetensors',
							    'text_encoder/pytorch_model.bin',
							    'text_encoder/model.safetensors',
							    'unet/diffusion_pytorch_model.bin',
							    # Removed: 'unet/diffusion_pytorch_model.safetensors',
							]
							self.assertFalse(is_safetensors_compatible(_lowercase       )       )
	def        lowerCamelCase__		(							self       )						->  Any:
							lowercase_      :   Dict        =			[
							    'text_encoder/pytorch_model.bin',
							    'text_encoder/model.safetensors',
							]
							self.assertTrue(is_safetensors_compatible(_lowercase       )       )
	def        lowerCamelCase__		(							self       )						->  Optional[Any]:
							lowercase_      :   List[Any]        =			[
							    'safety_checker/pytorch_model.bin',
							    'safety_checker/model.safetensors',
							    'vae/diffusion_pytorch_model.bin',
							    'vae/diffusion_pytorch_model.safetensors',
							    'text_encoder/pytorch_model.bin',
							    # Removed: 'text_encoder/model.safetensors',
							    'unet/diffusion_pytorch_model.bin',
							    'unet/diffusion_pytorch_model.safetensors',
							]
							self.assertFalse(is_safetensors_compatible(_lowercase       )       )
	def        lowerCamelCase__		(							self       )						->  Optional[int]:
							lowercase_      :   str        =			[
							    'safety_checker/pytorch_model.fp16.bin',
							    'safety_checker/model.fp16.safetensors',
							    'vae/diffusion_pytorch_model.fp16.bin',
							    'vae/diffusion_pytorch_model.fp16.safetensors',
							    'text_encoder/pytorch_model.fp16.bin',
							    'text_encoder/model.fp16.safetensors',
							    'unet/diffusion_pytorch_model.fp16.bin',
							    'unet/diffusion_pytorch_model.fp16.safetensors',
							]
							lowercase_      :   Union[str, Any]        =			'fp16'
							self.assertTrue(is_safetensors_compatible(_lowercase  ,      variant=_lowercase       )       )
	def        lowerCamelCase__		(							self       )						->  int:
							lowercase_      :   Optional[int]        =			[
							    'unet/diffusion_pytorch_model.fp16.bin',
							    'unet/diffusion_pytorch_model.fp16.safetensors',
							]
							lowercase_      :   Dict        =			'fp16'
							self.assertTrue(is_safetensors_compatible(_lowercase  ,      variant=_lowercase       )       )
	def        lowerCamelCase__		(							self       )						->  int:
							# pass variant but use the non-variant filenames
							lowercase_      :   Optional[int]        =			[
							    'unet/diffusion_pytorch_model.bin',
							    'unet/diffusion_pytorch_model.safetensors',
							]
							lowercase_      :   Dict        =			'fp16'
							self.assertTrue(is_safetensors_compatible(_lowercase  ,      variant=_lowercase       )       )
	def        lowerCamelCase__		(							self       )						->  Union[str, Any]:
							lowercase_      :   int        =			[
							    'safety_checker/pytorch_model.fp16.bin',
							    'safety_checker/model.fp16.safetensors',
							    'vae/diffusion_pytorch_model.fp16.bin',
							    'vae/diffusion_pytorch_model.fp16.safetensors',
							    'text_encoder/pytorch_model.fp16.bin',
							    'text_encoder/model.fp16.safetensors',
							    'unet/diffusion_pytorch_model.fp16.bin',
							    # Removed: 'unet/diffusion_pytorch_model.fp16.safetensors',
							]
							lowercase_      :   Dict        =			'fp16'
							self.assertFalse(is_safetensors_compatible(_lowercase  ,      variant=_lowercase       )       )
	def        lowerCamelCase__		(							self       )						->  int:
							lowercase_      :   str        =			[
							    'text_encoder/pytorch_model.fp16.bin',
							    'text_encoder/model.fp16.safetensors',
							]
							lowercase_      :   str        =			'fp16'
							self.assertTrue(is_safetensors_compatible(_lowercase  ,      variant=_lowercase       )       )
	def        lowerCamelCase__		(							self       )						->  List[str]:
							# pass variant but use the non-variant filenames
							lowercase_      :   List[Any]        =			[
							    'text_encoder/pytorch_model.bin',
							    'text_encoder/model.safetensors',
							]
							lowercase_      :   Dict        =			'fp16'
							self.assertTrue(is_safetensors_compatible(_lowercase  ,      variant=_lowercase       )       )







	def        lowerCamelCase__		(							self       )						->  List[str]:
							lowercase_      :   Union[str, Any]        =			[
							    'safety_checker/pytorch_model.fp16.bin',
							    'safety_checker/model.fp16.safetensors',
							    'vae/diffusion_pytorch_model.fp16.bin',
							    'vae/diffusion_pytorch_model.fp16.safetensors',
							    'text_encoder/pytorch_model.fp16.bin',
							    # 'text_encoder/model.fp16.safetensors',
							    'unet/diffusion_pytorch_model.fp16.bin',
							    'unet/diffusion_pytorch_model.fp16.safetensors',
							]
							lowercase_      :   Dict        =			'fp16'
							self.assertFalse(is_safetensors_compatible(_lowercase  ,      variant=_lowercase       )       )
