


"""simple docstring"""
import tempfile

import torch

from diffusers import PNDMScheduler

from .test_schedulers import SchedulerCommonTest



class 			_UpperCAmelCase (							__snake_case				):






		'''simple docstring'''

		lowerCamelCase__				=(PNDMScheduler,)
		lowerCamelCase__				=(('num_inference_steps', 50),)



		def 				SCREAMING_SNAKE_CASE  (self				,       **a_					):





			'''simple docstring'''

			__snake_case	:						Tuple						=				{
			    '''num_train_timesteps''': 10_00,
			    '''beta_start''': 0.0001,
			    '''beta_end''': 0.02,
			    '''beta_schedule''': '''linear''',
			}

			config.update(**a_					)
			return config



		def 				SCREAMING_SNAKE_CASE  (self				,       a_=0				,       **a_					):





			'''simple docstring'''

			__snake_case	:						Dict						=				dict(self.forward_default_kwargs					)
			__snake_case	:						str						=				kwargs.pop('''num_inference_steps'''				,       a_					)
			__snake_case	:						str						=				self.dummy_sample
			__snake_case	:						Any						=				0.1 * sample
			__snake_case	:						List[Any]						=				[residual + 0.2, residual + 0.15, residual + 0.1, residual + 0.05]

			for scheduler_class in self.scheduler_classes:
				__snake_case	:						Any						=				self.get_scheduler_config(**a_					)
				__snake_case	:						int						=				scheduler_class(**a_					)
				scheduler.set_timesteps(a_					)
				# copy over dummy past residuals
				__snake_case	:						Tuple						=				dummy_past_residuals[:]

				with tempfile.TemporaryDirectory() as tmpdirname:
					scheduler.save_config(a_					)
					__snake_case	:						List[Any]						=				scheduler_class.from_pretrained(a_					)
					new_scheduler.set_timesteps(a_					)
					# copy over dummy past residuals
					__snake_case	:						Any						=				dummy_past_residuals[:]

				__snake_case	:						int						=				scheduler.step_prk(a_				,       a_				,       a_				,       **a_					).prev_sample
				__snake_case	:						Union[str, Any]						=				new_scheduler.step_prk(a_				,       a_				,       a_				,       **a_					).prev_sample

				assert torch.sum(torch.abs(output - new_output					)					) < 1E-5, "Scheduler outputs are not identical"

				__snake_case	:						Tuple						=				scheduler.step_plms(a_				,       a_				,       a_				,       **a_					).prev_sample
				__snake_case	:						Tuple						=				new_scheduler.step_plms(a_				,       a_				,       a_				,       **a_					).prev_sample

				assert torch.sum(torch.abs(output - new_output					)					) < 1E-5, "Scheduler outputs are not identical"



		def 				SCREAMING_SNAKE_CASE  (self					):





			'''simple docstring'''

			pass



		def 				SCREAMING_SNAKE_CASE  (self				,       a_=0				,       **a_					):





			'''simple docstring'''

			__snake_case	:						int						=				dict(self.forward_default_kwargs					)
			__snake_case	:						Optional[Any]						=				kwargs.pop('''num_inference_steps'''				,       a_					)
			__snake_case	:						List[Any]						=				self.dummy_sample
			__snake_case	:						List[Any]						=				0.1 * sample
			__snake_case	:						List[str]						=				[residual + 0.2, residual + 0.15, residual + 0.1, residual + 0.05]

			for scheduler_class in self.scheduler_classes:
				__snake_case	:						str						=				self.get_scheduler_config()
				__snake_case	:						int						=				scheduler_class(**a_					)
				scheduler.set_timesteps(a_					)

				# copy over dummy past residuals (must be after setting timesteps)
				__snake_case	:						List[str]						=				dummy_past_residuals[:]

				with tempfile.TemporaryDirectory() as tmpdirname:
					scheduler.save_config(a_					)
					__snake_case	:						Any						=				scheduler_class.from_pretrained(a_					)
					# copy over dummy past residuals
					new_scheduler.set_timesteps(a_					)

					# copy over dummy past residual (must be after setting timesteps)
					__snake_case	:						str						=				dummy_past_residuals[:]

				__snake_case	:						Any						=				scheduler.step_prk(a_				,       a_				,       a_				,       **a_					).prev_sample
				__snake_case	:						Dict						=				new_scheduler.step_prk(a_				,       a_				,       a_				,       **a_					).prev_sample

				assert torch.sum(torch.abs(output - new_output					)					) < 1E-5, "Scheduler outputs are not identical"

				__snake_case	:						List[Any]						=				scheduler.step_plms(a_				,       a_				,       a_				,       **a_					).prev_sample
				__snake_case	:						Union[str, Any]						=				new_scheduler.step_plms(a_				,       a_				,       a_				,       **a_					).prev_sample

				assert torch.sum(torch.abs(output - new_output					)					) < 1E-5, "Scheduler outputs are not identical"



		def 				SCREAMING_SNAKE_CASE  (self				,       **a_					):





			'''simple docstring'''

			__snake_case	:						List[str]						=				self.scheduler_classes[0]
			__snake_case	:						List[Any]						=				self.get_scheduler_config(**a_					)
			__snake_case	:						Tuple						=				scheduler_class(**a_					)

			__snake_case	:						List[Any]						=				10
			__snake_case	:						List[str]						=				self.dummy_model()
			__snake_case	:						Any						=				self.dummy_sample_deter
			scheduler.set_timesteps(a_					)

			for i, t in enumerate(scheduler.prk_timesteps					):
				__snake_case	:						int						=				model(a_				,       a_					)
				__snake_case	:						Optional[int]						=				scheduler.step_prk(a_				,       a_				,       a_					).prev_sample

			for i, t in enumerate(scheduler.plms_timesteps					):
				__snake_case	:						int						=				model(a_				,       a_					)
				__snake_case	:						Optional[int]						=				scheduler.step_plms(a_				,       a_				,       a_					).prev_sample

			return sample



		def 				SCREAMING_SNAKE_CASE  (self					):





			'''simple docstring'''

			__snake_case	:						List[Any]						=				dict(self.forward_default_kwargs					)

			__snake_case	:						Any						=				kwargs.pop('''num_inference_steps'''				,       a_					)

			for scheduler_class in self.scheduler_classes:
				__snake_case	:						Optional[Any]						=				self.get_scheduler_config()
				__snake_case	:						Dict						=				scheduler_class(**a_					)

				__snake_case	:						List[Any]						=				self.dummy_sample
				__snake_case	:						List[Any]						=				0.1 * sample

				if num_inference_steps is not None and hasattr(a_				,       '''set_timesteps'''					):
					scheduler.set_timesteps(a_					)
				elif num_inference_steps is not None and not hasattr(a_				,       '''set_timesteps'''					):
					__snake_case	:						List[str]						=				num_inference_steps

				# copy over dummy past residuals (must be done after set_timesteps)
				__snake_case	:						Union[str, Any]						=				[residual + 0.2, residual + 0.15, residual + 0.1, residual + 0.05]
				__snake_case	:						List[Any]						=				dummy_past_residuals[:]

				__snake_case	:						str						=				scheduler.step_prk(a_				,       0				,       a_				,       **a_					).prev_sample
				__snake_case	:						Dict						=				scheduler.step_prk(a_				,       1				,       a_				,       **a_					).prev_sample

				self.assertEqual(output_a.shape				,       sample.shape					)
				self.assertEqual(output_a.shape				,       output_a.shape					)

				__snake_case	:						Tuple						=				scheduler.step_plms(a_				,       0				,       a_				,       **a_					).prev_sample
				__snake_case	:						List[str]						=				scheduler.step_plms(a_				,       1				,       a_				,       **a_					).prev_sample

				self.assertEqual(output_a.shape				,       sample.shape					)
				self.assertEqual(output_a.shape				,       output_a.shape					)



		def 				SCREAMING_SNAKE_CASE  (self					):





			'''simple docstring'''

			for timesteps in [1_00, 10_00]:
				self.check_over_configs(num_train_timesteps=a_					)



		def 				SCREAMING_SNAKE_CASE  (self					):





			'''simple docstring'''

			for steps_offset in [0, 1]:
				self.check_over_configs(steps_offset=a_					)

			__snake_case	:						List[str]						=				self.scheduler_classes[0]
			__snake_case	:						Tuple						=				self.get_scheduler_config(steps_offset=1					)
			__snake_case	:						int						=				scheduler_class(**a_					)
			scheduler.set_timesteps(10					)
			assert torch.equal(
			    scheduler.timesteps				,       torch.LongTensor(
			        [9_01, 8_51, 8_51, 8_01, 8_01, 7_51, 7_51, 7_01, 7_01, 6_51, 6_51, 6_01, 6_01, 5_01, 4_01, 3_01, 2_01, 1_01, 1]					)				,       )



		def 				SCREAMING_SNAKE_CASE  (self					):





			'''simple docstring'''

			for beta_start, beta_end in zip([0.0001, 0.001]				,       [0.002, 0.02]					):
				self.check_over_configs(beta_start=a_				,       beta_end=a_					)



		def 				SCREAMING_SNAKE_CASE  (self					):





			'''simple docstring'''

			for schedule in ["linear", "squaredcos_cap_v2"]:
				self.check_over_configs(beta_schedule=a_					)



		def 				SCREAMING_SNAKE_CASE  (self					):





			'''simple docstring'''

			for prediction_type in ["epsilon", "v_prediction"]:
				self.check_over_configs(prediction_type=a_					)



		def 				SCREAMING_SNAKE_CASE  (self					):





			'''simple docstring'''

			for t in [1, 5, 10]:
				self.check_over_forward(time_step=a_					)



		def 				SCREAMING_SNAKE_CASE  (self					):





			'''simple docstring'''

			for t, num_inference_steps in zip([1, 5, 10]				,       [10, 50, 1_00]					):
				self.check_over_forward(num_inference_steps=a_					)



		def 				SCREAMING_SNAKE_CASE  (self					):





			'''simple docstring'''

			__snake_case	:						Union[str, Any]						=				27

			for scheduler_class in self.scheduler_classes:
				__snake_case	:						List[str]						=				self.dummy_sample
				__snake_case	:						Dict						=				0.1 * sample

				__snake_case	:						Tuple						=				self.get_scheduler_config()
				__snake_case	:						List[str]						=				scheduler_class(**a_					)

				scheduler.set_timesteps(a_					)

				# before power of 3 fix, would error on first step, so we only need to do two
				for i, t in enumerate(scheduler.prk_timesteps[:2]					):
					__snake_case	:						Optional[Any]						=				scheduler.step_prk(a_				,       a_				,       a_					).prev_sample



		def 				SCREAMING_SNAKE_CASE  (self					):





			'''simple docstring'''

			with self.assertRaises(a_					):
				__snake_case	:						List[Any]						=				self.scheduler_classes[0]
				__snake_case	:						Optional[int]						=				self.get_scheduler_config()
				__snake_case	:						Tuple						=				scheduler_class(**a_					)

				scheduler.step_plms(self.dummy_sample				,       1				,       self.dummy_sample					).prev_sample



		def 				SCREAMING_SNAKE_CASE  (self					):





			'''simple docstring'''

			__snake_case	:						str						=				self.full_loop()
			__snake_case	:						Any						=				torch.sum(torch.abs(a_					)					)
			__snake_case	:						Dict						=				torch.mean(torch.abs(a_					)					)

			assert abs(result_sum.item() - 198.1318					) < 1E-2
			assert abs(result_mean.item() - 0.2580					) < 1E-3



		def 				SCREAMING_SNAKE_CASE  (self					):





			'''simple docstring'''

			__snake_case	:						int						=				self.full_loop(prediction_type='''v_prediction'''					)
			__snake_case	:						str						=				torch.sum(torch.abs(a_					)					)
			__snake_case	:						int						=				torch.mean(torch.abs(a_					)					)

			assert abs(result_sum.item() - 67.3986					) < 1E-2
			assert abs(result_mean.item() - 0.0878					) < 1E-3



		def 				SCREAMING_SNAKE_CASE  (self					):





			'''simple docstring'''

			__snake_case	:						int						=				self.full_loop(set_alpha_to_one=a_				,       beta_start=0.01					)
			__snake_case	:						Any						=				torch.sum(torch.abs(a_					)					)
			__snake_case	:						int						=				torch.mean(torch.abs(a_					)					)

			assert abs(result_sum.item() - 230.0399					) < 1E-2
			assert abs(result_mean.item() - 0.2995					) < 1E-3



		def 				SCREAMING_SNAKE_CASE  (self					):





			'''simple docstring'''

			__snake_case	:						Optional[Any]						=				self.full_loop(set_alpha_to_one=a_				,       beta_start=0.01					)
			__snake_case	:						Any						=				torch.sum(torch.abs(a_					)					)
			__snake_case	:						List[str]						=				torch.mean(torch.abs(a_					)					)

			assert abs(result_sum.item() - 186.9482					) < 1E-2
			assert abs(result_mean.item() - 0.2434					) < 1E-3









"""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 (							__snake_case,				unittest.TestCase				):






		'''simple docstring'''

		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 				SCREAMING_SNAKE_CASE  (self					):





			'''simple docstring'''

			return 32



		@property
		def 				SCREAMING_SNAKE_CASE  (self					):





			'''simple docstring'''

			return 32



		@property
		def 				SCREAMING_SNAKE_CASE  (self					):





			'''simple docstring'''

			return self.time_input_dim * 4



		@property
		def 				SCREAMING_SNAKE_CASE  (self					):





			'''simple docstring'''

			return 8



		@property
		def 				SCREAMING_SNAKE_CASE  (self					):





			'''simple docstring'''

			__snake_case	:						Optional[Any]						=				CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''					)
			return tokenizer



		@property
		def 				SCREAMING_SNAKE_CASE  (self					):





			'''simple docstring'''

			torch.manual_seed(0					)
			__snake_case	:						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-05				,       num_attention_heads=4				,       num_hidden_layers=5				,       pad_token_id=1				,       vocab_size=10_00				,       )
			return CLIPTextModelWithProjection(a_					)



		@property
		def 				SCREAMING_SNAKE_CASE  (self					):





			'''simple docstring'''

			torch.manual_seed(0					)

			__snake_case	:						Optional[int]						=				{
			    '''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,
			}

			__snake_case	:						Optional[int]						=				PriorTransformer(**a_					)
			return model



		@property
		def 				SCREAMING_SNAKE_CASE  (self					):





			'''simple docstring'''

			torch.manual_seed(0					)

			__snake_case	:						Optional[int]						=				{
			    '''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,
			    ),
			}
			__snake_case	:						List[Any]						=				ShapERenderer(**a_					)
			return model



		def 				SCREAMING_SNAKE_CASE  (self					):





			'''simple docstring'''

			__snake_case	:						int						=				self.dummy_prior
			__snake_case	:						str						=				self.dummy_text_encoder
			__snake_case	:						str						=				self.dummy_tokenizer
			__snake_case	:						Tuple						=				self.dummy_renderer

			__snake_case	:						int						=				HeunDiscreteScheduler(
			    beta_schedule='''exp'''				,       num_train_timesteps=10_24				,       prediction_type='''sample'''				,       use_karras_sigmas=a_				,       clip_sample=a_				,       clip_sample_range=1.0				,       )
			__snake_case	:						Union[str, Any]						=				{
			    '''prior''': prior,
			    '''text_encoder''': text_encoder,
			    '''tokenizer''': tokenizer,
			    '''renderer''': renderer,
			    '''scheduler''': scheduler,
			}

			return components



		def 				SCREAMING_SNAKE_CASE  (self				,       a_				,       a_=0					):





			'''simple docstring'''

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



		def 				SCREAMING_SNAKE_CASE  (self					):





			'''simple docstring'''

			__snake_case	:						Optional[int]						=				'''cpu'''

			__snake_case	:						str						=				self.get_dummy_components()

			__snake_case	:						List[Any]						=				self.pipeline_class(**a_					)
			__snake_case	:						str						=				pipe.to(a_					)

			pipe.set_progress_bar_config(disable=a_					)

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

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

			__snake_case	:						List[Any]						=				np.array(
			    [
			        0.0003_9216,
			        0.0003_9216,
			        0.0003_9216,
			        0.0003_9216,
			        0.0003_9216,
			        0.0003_9216,
			        0.0003_9216,
			        0.0003_9216,
			        0.0003_9216,
			    ]					)

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



		def 				SCREAMING_SNAKE_CASE  (self					):





			'''simple docstring'''

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



		def 				SCREAMING_SNAKE_CASE  (self					):





			'''simple docstring'''

			__snake_case	:						Any						=				torch_device == '''cpu'''
			__snake_case	:						Tuple						=				True

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



		def 				SCREAMING_SNAKE_CASE  (self					):





			'''simple docstring'''

			__snake_case	:						int						=				self.get_dummy_components()
			__snake_case	:						int						=				self.pipeline_class(**a_					)
			__snake_case	:						int						=				pipe.to(a_					)
			pipe.set_progress_bar_config(disable=a_					)

			__snake_case	:						Optional[Any]						=				1
			__snake_case	:						List[Any]						=				2

			__snake_case	:						int						=				self.get_dummy_inputs(a_					)

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

			__snake_case	:						Union[str, Any]						=				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				):






		'''simple docstring'''

		def 				SCREAMING_SNAKE_CASE  (self					):





			'''simple docstring'''

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



		def 				SCREAMING_SNAKE_CASE  (self					):





			'''simple docstring'''

			__snake_case	:						List[str]						=				load_numpy(
			    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
			    '''/shap_e/test_shap_e_np_out.npy'''					)
			__snake_case	:						Optional[Any]						=				ShapEPipeline.from_pretrained('''openai/shap-e'''					)
			__snake_case	:						Optional[int]						=				pipe.to(a_					)
			pipe.set_progress_bar_config(disable=a_					)

			__snake_case	:						Optional[Any]						=				torch.Generator(device=a_					).manual_seed(0					)

			__snake_case	:						str						=				pipe(
			    '''a shark'''				,       generator=a_				,       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(a_				,       a_					)





