






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 				lowerCAmelCase_	(		__snake_case   ,		unittest.TestCase			):
				_UpperCamelCase					:       str							       = ShapEPipeline
				_UpperCamelCase					:       Any							       = ["prompt"]
				_UpperCamelCase					:       int							       = ["prompt"]
				_UpperCamelCase					:       Union[str, Any]							       = [
				    "num_images_per_prompt",
				    "num_inference_steps",
				    "generator",
				    "latents",
				    "guidance_scale",
				    "frame_size",
				    "output_type",
				    "return_dict",
				]
				_UpperCamelCase					:       Optional[Any]							       = False


				@property
				def 						__a					(					self      ):
											return 3_2


				@property
				def 						__a					(					self      ):
											return 3_2


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


				@property
				def 						__a					(					self      ):
											return 8


				@property
				def 						__a					(					self      ):
											_lowercase							:   Any 							=   CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'      )
											return tokenizer


				@property
				def 						__a					(					self      ):
											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=3_7 ,						layer_norm_eps=1E-05 ,						num_attention_heads=4 ,						num_hidden_layers=5 ,						pad_token_id=1 ,						vocab_size=1_0_0_0 ,						)
											return CLIPTextModelWithProjection(_lowerCAmelCase      )


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

											_lowercase							:   Optional[int] 							=   {
											    'num_attention_heads': 2,
											    'attention_head_dim': 1_6,
											    'embedding_dim': self.time_input_dim,
											    'num_embeddings': 3_2,
											    '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							:   Optional[Any] 							=   PriorTransformer(**_lowerCAmelCase      )
											return model


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

											_lowercase							:   Optional[int] 							=   {
											    'param_shapes': (
											        (self.renderer_dim, 9_3),
											        (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': 1_2,
											    'background': (
											        0.1,
											        0.1,
											        0.1,
											    ),
											}
											_lowercase							:   List[Any] 							=   ShapERenderer(**_lowerCAmelCase      )
											return model


				def 						__a					(					self      ):
											_lowercase							:   Optional[Any] 							=   self.dummy_prior
											_lowercase							:   Dict 							=   self.dummy_text_encoder
											_lowercase							:   List[str] 							=   self.dummy_tokenizer
											_lowercase							:   Union[str, Any] 							=   self.dummy_renderer

											_lowercase							:   List[str] 							=   HeunDiscreteScheduler(
											    beta_schedule='exp' ,						num_train_timesteps=1_0_2_4 ,						prediction_type='sample' ,						use_karras_sigmas=_lowerCAmelCase ,						clip_sample=_lowerCAmelCase ,						clip_sample_range=1.0 ,						)
											_lowercase							:   List[str] 							=   {
											    'prior': prior,
											    'text_encoder': text_encoder,
											    'tokenizer': tokenizer,
											    'renderer': renderer,
											    'scheduler': scheduler,
											}

											return components


				def 						__a					(					self ,						_lowerCAmelCase ,						_lowerCAmelCase=0      ):
											if str(_lowerCAmelCase      ).startswith('mps'      ):
																		_lowercase							:   Optional[Any] 							=   torch.manual_seed(_lowerCAmelCase      )
											else:
																		_lowercase							:   Optional[int] 							=   torch.Generator(device=_lowerCAmelCase      ).manual_seed(_lowerCAmelCase      )
											_lowercase							:   List[Any] 							=   {
											    'prompt': 'horse',
											    'generator': generator,
											    'num_inference_steps': 1,
											    'frame_size': 3_2,
											    'output_type': 'np',
											}
											return inputs


				def 						__a					(					self      ):
											_lowercase							:   Optional[int] 							=   'cpu'

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

											_lowercase							:   Tuple 							=   self.pipeline_class(**_lowerCAmelCase      )
											_lowercase							:   Union[str, Any] 							=   pipe.to(_lowerCAmelCase      )

											pipe.set_progress_bar_config(disable=_lowerCAmelCase      )

											_lowercase							:   Union[str, Any] 							=   pipe(**self.get_dummy_inputs(_lowerCAmelCase      )      )
											_lowercase							:   str 							=   output.images[0]
											_lowercase							:   Optional[Any] 							=   image[0, -3:, -3:, -1]

											assert image.shape == (2_0, 3_2, 3_2, 3)

											_lowercase							:   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 						__a					(					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 						__a					(					self      ):
											_lowercase							:   List[Any] 							=   torch_device == 'cpu'
											_lowercase							:   Any 							=   True

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


				def 						__a					(					self      ):
											_lowercase							:   Union[str, Any] 							=   self.get_dummy_components()
											_lowercase							:   Optional[int] 							=   self.pipeline_class(**_lowerCAmelCase      )
											_lowercase							:   Any 							=   pipe.to(_lowerCAmelCase      )
											pipe.set_progress_bar_config(disable=_lowerCAmelCase      )

											_lowercase							:   str 							=   1
											_lowercase							:   Optional[int] 							=   2

											_lowercase							:   List[str] 							=   self.get_dummy_inputs(_lowerCAmelCase      )

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

											_lowercase							:   Optional[int] 							=   pipe(**_lowerCAmelCase ,						num_images_per_prompt=_lowerCAmelCase      )[0]

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




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


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


				def 						__a					(					self      ):
											_lowercase							:   List[str] 							=   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[str] 							=   pipe.to(_lowerCAmelCase      )
											pipe.set_progress_bar_config(disable=_lowerCAmelCase      )

											_lowercase							:   Tuple 							=   torch.Generator(device=_lowerCAmelCase      ).manual_seed(0      )

											_lowercase							:   int 							=   pipe(
											    'a shark' ,						generator=_lowerCAmelCase ,						guidance_scale=15.0 ,						num_inference_steps=6_4 ,						frame_size=6_4 ,						output_type='np' ,						).images[0]

											assert images.shape == (2_0, 6_4, 6_4, 3)

											assert_mean_pixel_difference(_lowerCAmelCase ,						_lowerCAmelCase      )









"""simple docstring"""

import gc
import tempfile
import unittest

import numpy as np
import torch

from diffusers import VersatileDiffusionTextToImagePipeline
from diffusers.utils.testing_utils import nightly, require_torch_gpu, torch_device


A_						:    Union[str, Any]				=      False
class     lowerCAmelCase__						(	unittest.TestCase ):


	'''simple docstring'''







	pass




@nightly
@require_torch_gpu
class     lowerCAmelCase__						(	unittest.TestCase ):


	'''simple docstring'''







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





					"""simple docstring"""


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





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





					"""simple docstring"""


					SCREAMING_SNAKE_CASE						:							Tuple							=     VersatileDiffusionTextToImagePipeline.from_pretrained('shi-labs/versatile-diffusion'    )
					# remove text_unet
					pipe.remove_unused_weights()
					pipe.to(_SCREAMING_SNAKE_CASE    )
					pipe.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE    )

					SCREAMING_SNAKE_CASE						:							Union[str, Any]							=     'A painting of a squirrel eating a burger '
					SCREAMING_SNAKE_CASE						:							List[str]							=     torch.manual_seed(0    )
					SCREAMING_SNAKE_CASE						:							int							=     pipe(
					    prompt=_SCREAMING_SNAKE_CASE     ,       generator=_SCREAMING_SNAKE_CASE     ,       guidance_scale=7.5     ,       num_inference_steps=2     ,       output_type='numpy'    ).images

					with tempfile.TemporaryDirectory() as tmpdirname:
									pipe.save_pretrained(_SCREAMING_SNAKE_CASE    )
									SCREAMING_SNAKE_CASE						:							Dict							=     VersatileDiffusionTextToImagePipeline.from_pretrained(_SCREAMING_SNAKE_CASE    )
					pipe.to(_SCREAMING_SNAKE_CASE    )
					pipe.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE    )

					SCREAMING_SNAKE_CASE						:							Union[str, Any]							=     generator.manual_seed(0    )
					SCREAMING_SNAKE_CASE						:							Tuple							=     pipe(
					    prompt=_SCREAMING_SNAKE_CASE     ,       generator=_SCREAMING_SNAKE_CASE     ,       guidance_scale=7.5     ,       num_inference_steps=2     ,       output_type='numpy'    ).images

					assert np.abs(image - new_image    ).sum() < 1E-5, "Models don't have the same forward pass"







	def 			_lowerCAmelCase     (      self				:       Dict    )						->	Dict:





					"""simple docstring"""


					SCREAMING_SNAKE_CASE						:							List[Any]							=     VersatileDiffusionTextToImagePipeline.from_pretrained(
					    'shi-labs/versatile-diffusion'     ,       torch_dtype=torch.floataa    )
					pipe.to(_SCREAMING_SNAKE_CASE    )
					pipe.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE    )

					SCREAMING_SNAKE_CASE						:							Optional[int]							=     'A painting of a squirrel eating a burger '
					SCREAMING_SNAKE_CASE						:							Dict							=     torch.manual_seed(0    )
					SCREAMING_SNAKE_CASE						:							Any							=     pipe(
					    prompt=_SCREAMING_SNAKE_CASE     ,       generator=_SCREAMING_SNAKE_CASE     ,       guidance_scale=7.5     ,       num_inference_steps=50     ,       output_type='numpy'    ).images

					SCREAMING_SNAKE_CASE						:							Tuple							=     image[0, 253:256, 253:256, -1]

					assert image.shape == (1, 512, 512, 3)
					SCREAMING_SNAKE_CASE						:							Union[str, Any]							=     np.array([0.3_3_6_7, 0.3_1_6_9, 0.2_6_5_6, 0.3_8_7_0, 0.4_7_9_0, 0.3_7_9_6, 0.4_0_0_9, 0.4_8_7_8, 0.4_7_7_8]    )

					assert np.abs(image_slice.flatten() - expected_slice    ).max() < 1E-2





