





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_							(  UpperCamelCase__					,	unittest.TestCase     ):
							lowerCamelCase_							     =       ShapEPipeline
							lowerCamelCase_							     =       ["prompt"]
							lowerCamelCase_							     =       ["prompt"]
							lowerCamelCase_							     =       [
							    "num_images_per_prompt",
							    "num_inference_steps",
							    "generator",
							    "latents",
							    "guidance_scale",
							    "frame_size",
							    "output_type",
							    "return_dict",
							]
							lowerCamelCase_							     =       False


							@property
							def   _snake_case	(				self     :Union[str, Any]		) -> Any:



									"""simple docstring"""



									return 32


							@property
							def   _snake_case	(				self     :str		) -> List[Any]:



									"""simple docstring"""



									return 32


							@property
							def   _snake_case	(				self     :Tuple		) -> Dict:



									"""simple docstring"""



									return self.time_input_dim * 4


							@property
							def   _snake_case	(				self     :Optional[Any]		) -> List[str]:



									"""simple docstring"""



									return 8


							@property
							def   _snake_case	(				self     :Any		) -> int:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__								=   CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""		)
									return tokenizer


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



									"""simple docstring"""



									torch.manual_seed(0		)
									SCREAMING_SNAKE_CASE__								=   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=1000				,					)
									return CLIPTextModelWithProjection(__A		)


							@property
							def   _snake_case	(				self     :Tuple		) -> Any:



									"""simple docstring"""



									torch.manual_seed(0		)

									SCREAMING_SNAKE_CASE__								=   {
									    """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,
									}

									SCREAMING_SNAKE_CASE__								=   PriorTransformer(**__A		)
									return model


							@property
							def   _snake_case	(				self     :str		) -> int:



									"""simple docstring"""



									torch.manual_seed(0		)

									SCREAMING_SNAKE_CASE__								=   {
									    """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,
									    ),
									}
									SCREAMING_SNAKE_CASE__								=   ShapERenderer(**__A		)
									return model


							def   _snake_case	(				self     :int		) -> List[str]:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__								=   self.dummy_prior
									SCREAMING_SNAKE_CASE__								=   self.dummy_text_encoder
									SCREAMING_SNAKE_CASE__								=   self.dummy_tokenizer
									SCREAMING_SNAKE_CASE__								=   self.dummy_renderer

									SCREAMING_SNAKE_CASE__								=   HeunDiscreteScheduler(
									    beta_schedule="""exp"""				,					num_train_timesteps=1024				,					prediction_type="""sample"""				,					use_karras_sigmas=__A				,					clip_sample=__A				,					clip_sample_range=1.0				,					)
									SCREAMING_SNAKE_CASE__								=   {
									    """prior""": prior,
									    """text_encoder""": text_encoder,
									    """tokenizer""": tokenizer,
									    """renderer""": renderer,
									    """scheduler""": scheduler,
									}

									return components


							def   _snake_case	(				self     :List[str]				,					__A     :List[Any]				,					__A     :List[Any]=0		) -> str:



									"""simple docstring"""



									if str(__A		).startswith("""mps"""		):
											SCREAMING_SNAKE_CASE__								=   torch.manual_seed(__A		)
									else:
											SCREAMING_SNAKE_CASE__								=   torch.Generator(device=__A		).manual_seed(__A		)
									SCREAMING_SNAKE_CASE__								=   {
									    """prompt""": """horse""",
									    """generator""": generator,
									    """num_inference_steps""": 1,
									    """frame_size""": 32,
									    """output_type""": """np""",
									}
									return inputs


							def   _snake_case	(				self     :Dict		) -> List[str]:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__								=   """cpu"""

									SCREAMING_SNAKE_CASE__								=   self.get_dummy_components()

									SCREAMING_SNAKE_CASE__								=   self.pipeline_class(**__A		)
									SCREAMING_SNAKE_CASE__								=   pipe.to(__A		)

									pipe.set_progress_bar_config(disable=__A		)

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

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

									SCREAMING_SNAKE_CASE__								=   np.array(
									    [
									        0.0_0_0_3_9_2_1_6,
									        0.0_0_0_3_9_2_1_6,
									        0.0_0_0_3_9_2_1_6,
									        0.0_0_0_3_9_2_1_6,
									        0.0_0_0_3_9_2_1_6,
									        0.0_0_0_3_9_2_1_6,
									        0.0_0_0_3_9_2_1_6,
									        0.0_0_0_3_9_2_1_6,
									        0.0_0_0_3_9_2_1_6,
									    ]		)

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


							def   _snake_case	(				self     :Any		) -> Dict:



									"""simple docstring"""



									self._test_inference_batch_consistent(batch_sizes=[1, 2]		)


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



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__								=   torch_device == """cpu"""
									SCREAMING_SNAKE_CASE__								=   True

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


							def   _snake_case	(				self     :Union[str, Any]		) -> List[str]:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__								=   self.get_dummy_components()
									SCREAMING_SNAKE_CASE__								=   self.pipeline_class(**__A		)
									SCREAMING_SNAKE_CASE__								=   pipe.to(__A		)
									pipe.set_progress_bar_config(disable=__A		)

									SCREAMING_SNAKE_CASE__								=   1
									SCREAMING_SNAKE_CASE__								=   2

									SCREAMING_SNAKE_CASE__								=   self.get_dummy_inputs(__A		)

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

									SCREAMING_SNAKE_CASE__								=   pipe(**__A				,					num_images_per_prompt=__A		)[0]

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




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


							def   _snake_case	(				self     :Tuple		) -> Optional[int]:



									"""simple docstring"""



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


							def   _snake_case	(				self     :str		) -> int:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__								=   load_numpy(
									    """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"""
									    """/shap_e/test_shap_e_np_out.npy"""		)
									SCREAMING_SNAKE_CASE__								=   ShapEPipeline.from_pretrained("""openai/shap-e"""		)
									SCREAMING_SNAKE_CASE__								=   pipe.to(__A		)
									pipe.set_progress_bar_config(disable=__A		)

									SCREAMING_SNAKE_CASE__								=   torch.Generator(device=__A		).manual_seed(0		)

									SCREAMING_SNAKE_CASE__								=   pipe(
									    """a shark"""				,					generator=__A				,					guidance_scale=1_5.0				,					num_inference_steps=64				,					frame_size=64				,					output_type="""np"""				,					).images[0]

									assert images.shape == (20, 64, 64, 3)

									assert_mean_pixel_difference(__A				,					__A		)






import gc
import unittest

import numpy as np
import torch

from diffusers import StableDiffusionKDiffusionPipeline
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu


enable_full_determinism()




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


							def   _snake_case	(				self     :Union[str, Any]		) -> List[str]:



									"""simple docstring"""



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


							def   _snake_case	(				self     :Any		) -> Optional[int]:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__								=   StableDiffusionKDiffusionPipeline.from_pretrained("""CompVis/stable-diffusion-v1-4"""		)
									SCREAMING_SNAKE_CASE__								=   sd_pipe.to(__A		)
									sd_pipe.set_progress_bar_config(disable=__A		)

									sd_pipe.set_scheduler("""sample_euler"""		)

									SCREAMING_SNAKE_CASE__								=   """A painting of a squirrel eating a burger"""
									SCREAMING_SNAKE_CASE__								=   torch.manual_seed(0		)
									SCREAMING_SNAKE_CASE__								=   sd_pipe([prompt]				,					generator=__A				,					guidance_scale=9.0				,					num_inference_steps=20				,					output_type="""np"""		)

									SCREAMING_SNAKE_CASE__								=   output.images

									SCREAMING_SNAKE_CASE__								=   image[0, -3:, -3:, -1]

									assert image.shape == (1, 512, 512, 3)
									SCREAMING_SNAKE_CASE__								=   np.array([0.0_4_4_7, 0.0_4_9_2, 0.0_4_6_8, 0.0_4_0_8, 0.0_3_8_3, 0.0_4_0_8, 0.0_3_5_4, 0.0_3_8_0, 0.0_3_3_9]		)

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


							def   _snake_case	(				self     :str		) -> Union[str, Any]:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__								=   StableDiffusionKDiffusionPipeline.from_pretrained("""stabilityai/stable-diffusion-2-1-base"""		)
									SCREAMING_SNAKE_CASE__								=   sd_pipe.to(__A		)
									sd_pipe.set_progress_bar_config(disable=__A		)

									sd_pipe.set_scheduler("""sample_euler"""		)

									SCREAMING_SNAKE_CASE__								=   """A painting of a squirrel eating a burger"""
									SCREAMING_SNAKE_CASE__								=   torch.manual_seed(0		)
									SCREAMING_SNAKE_CASE__								=   sd_pipe([prompt]				,					generator=__A				,					guidance_scale=9.0				,					num_inference_steps=20				,					output_type="""np"""		)

									SCREAMING_SNAKE_CASE__								=   output.images

									SCREAMING_SNAKE_CASE__								=   image[0, -3:, -3:, -1]

									assert image.shape == (1, 512, 512, 3)
									SCREAMING_SNAKE_CASE__								=   np.array([0.1_2_3_7, 0.1_3_2_0, 0.1_4_3_8, 0.1_3_5_9, 0.1_3_9_0, 0.1_1_3_2, 0.1_2_7_7, 0.1_1_7_5, 0.1_1_1_2]		)

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


							def   _snake_case	(				self     :Tuple		) -> Optional[Any]:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__								=   StableDiffusionKDiffusionPipeline.from_pretrained("""stabilityai/stable-diffusion-2-1-base"""		)
									SCREAMING_SNAKE_CASE__								=   sd_pipe.to(__A		)
									sd_pipe.set_progress_bar_config(disable=__A		)

									sd_pipe.set_scheduler("""sample_dpmpp_2m"""		)

									SCREAMING_SNAKE_CASE__								=   """A painting of a squirrel eating a burger"""
									SCREAMING_SNAKE_CASE__								=   torch.manual_seed(0		)
									SCREAMING_SNAKE_CASE__								=   sd_pipe(
									    [prompt]				,					generator=__A				,					guidance_scale=7.5				,					num_inference_steps=15				,					output_type="""np"""				,					use_karras_sigmas=__A				,					)

									SCREAMING_SNAKE_CASE__								=   output.images

									SCREAMING_SNAKE_CASE__								=   image[0, -3:, -3:, -1]

									assert image.shape == (1, 512, 512, 3)
									SCREAMING_SNAKE_CASE__								=   np.array(
									    [0.1_1_3_8_1_6_8_9, 0.1_2_1_1_2_9_2_1, 0.1_3_8_9_4_5_7, 0.1_2_5_4_9_6_0_6, 0.1_2_4_4_9_6_4, 0.1_0_8_3_1_5_1_7, 0.1_1_5_6_2_8_6_6, 0.1_0_8_6_7_8_1_6, 0.1_0_4_9_9_0_4_8]		)

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