


"""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		:       int							 =ShapEPipeline
							__UpperCAmelCase		:       Union[str, Any]							 =["""prompt"""]
							__UpperCAmelCase		:       int							 =["""prompt"""]
							__UpperCAmelCase		:       List[Any]							 =[
							    """num_images_per_prompt""",
							    """num_inference_steps""",
							    """generator""",
							    """latents""",
							    """guidance_scale""",
							    """frame_size""",
							    """output_type""",
							    """return_dict""",
							]
							__UpperCAmelCase		:       Dict							 =False

							@property
							def      snake_case							(	self      ):
										return 32

							@property
							def      snake_case							(	self      ):
										return 32

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

							@property
							def      snake_case							(	self      ):
										return 8

							@property
							def      snake_case							(	self      ):
										__lowerCAmelCase							       =     CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"      )
										return tokenizer

							@property
							def      snake_case							(	self      ):
										torch.manual_seed(0      )
										__lowerCAmelCase							       =     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=10_00       ,  )
										return CLIPTextModelWithProjection(__a      )

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

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

										__lowerCAmelCase							       =     PriorTransformer(**__a      )
										return model

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

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

							def      snake_case							(	self      ):
										__lowerCAmelCase							       =     self.dummy_prior
										__lowerCAmelCase							       =     self.dummy_text_encoder
										__lowerCAmelCase							       =     self.dummy_tokenizer
										__lowerCAmelCase							       =     self.dummy_renderer

										__lowerCAmelCase							       =     HeunDiscreteScheduler(
										    beta_schedule="exp"       ,  num_train_timesteps=10_24       ,  prediction_type="sample"       ,  use_karras_sigmas=__a       ,  clip_sample=__a       ,  clip_sample_range=1.0       ,  )
										__lowerCAmelCase							       =     {
										    "prior": prior,
										    "text_encoder": text_encoder,
										    "tokenizer": tokenizer,
										    "renderer": renderer,
										    "scheduler": scheduler,
										}

										return components

							def      snake_case							(	self       ,  __a       ,  __a=0      ):
										if str(__a      ).startswith("mps"      ):
													__lowerCAmelCase							       =     torch.manual_seed(__a      )
										else:
													__lowerCAmelCase							       =     torch.Generator(device=__a      ).manual_seed(__a      )
										__lowerCAmelCase							       =     {
										    "prompt": "horse",
										    "generator": generator,
										    "num_inference_steps": 1,
										    "frame_size": 32,
										    "output_type": "np",
										}
										return inputs

							def      snake_case							(	self      ):
										__lowerCAmelCase							       =     "cpu"

										__lowerCAmelCase							       =     self.get_dummy_components()

										__lowerCAmelCase							       =     self.pipeline_class(**__a      )
										__lowerCAmelCase							       =     pipe.to(__a      )

										pipe.set_progress_bar_config(disable=__a      )

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

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

										__lowerCAmelCase							       =     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      ):
										# 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      ):
										__lowerCAmelCase							       =     torch_device == "cpu"
										__lowerCAmelCase							       =     True

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



							def      snake_case							(	self      ):
										__lowerCAmelCase							       =     self.get_dummy_components()
										__lowerCAmelCase							       =     self.pipeline_class(**__a      )
										__lowerCAmelCase							       =     pipe.to(__a      )
										pipe.set_progress_bar_config(disable=__a      )

										__lowerCAmelCase							       =     1
										__lowerCAmelCase							       =     2

										__lowerCAmelCase							       =     self.get_dummy_inputs(__a      )

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

										__lowerCAmelCase							       =     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      snake_case							(	self      ):
										# clean up the VRAM after each test
										super().tearDown()
										gc.collect()
										torch.cuda.empty_cache()



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

										__lowerCAmelCase							       =     torch.Generator(device=__a      ).manual_seed(0      )

										__lowerCAmelCase							       =     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      )








"""simple docstring"""


import torch

from diffusers import KDPMaDiscreteScheduler
from diffusers.utils import torch_device

from .test_schedulers import SchedulerCommonTest
class     _UpperCamelCase    (  lowerCAmelCase__						):
							'''simple docstring'''







							__UpperCAmelCase		:       Tuple							 =(KDPMaDiscreteScheduler,)
							__UpperCAmelCase		:       Optional[Any]							 =1_0

							def      snake_case							(	self       ,  **__a      ):
										__lowerCAmelCase							       =     {
										    "num_train_timesteps": 11_00,
										    "beta_start": 0.0_0_0_1,
										    "beta_end": 0.0_2,
										    "beta_schedule": "linear",
										}

										config.update(**__a      )
										return config

							def      snake_case							(	self      ):
										for timesteps in [10, 50, 1_00, 10_00]:
													self.check_over_configs(num_train_timesteps=__a      )

							def      snake_case							(	self      ):
										for beta_start, beta_end in zip([0.0_0_0_0_1, 0.0_0_0_1, 0.0_0_1]       ,  [0.0_0_0_2, 0.0_0_2, 0.0_2]      ):
													self.check_over_configs(beta_start=__a       ,  beta_end=__a      )

							def      snake_case							(	self      ):
										for schedule in ["linear", "scaled_linear"]:
													self.check_over_configs(beta_schedule=__a      )

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

							def      snake_case							(	self      ):
										__lowerCAmelCase							       =     self.scheduler_classes[0]
										__lowerCAmelCase							       =     self.get_scheduler_config(prediction_type="v_prediction"      )
										__lowerCAmelCase							       =     scheduler_class(**__a      )

										scheduler.set_timesteps(self.num_inference_steps      )

										__lowerCAmelCase							       =     self.dummy_model()
										__lowerCAmelCase							       =     self.dummy_sample_deter * scheduler.init_noise_sigma
										__lowerCAmelCase							       =     sample.to(__a      )

										for i, t in enumerate(scheduler.timesteps      ):
													__lowerCAmelCase							       =     scheduler.scale_model_input(__a       ,  __a      )

													__lowerCAmelCase							       =     model(__a       ,  __a      )

													__lowerCAmelCase							       =     scheduler.step(__a       ,  __a       ,  __a      )
													__lowerCAmelCase							       =     output.prev_sample

										__lowerCAmelCase							       =     torch.sum(torch.abs(__a      )      )
										__lowerCAmelCase							       =     torch.mean(torch.abs(__a      )      )

										if torch_device in ["cpu", "mps"]:
													assert abs(result_sum.item() - 4.6_9_3_4e-0_7      ) < 1e-2
													assert abs(result_mean.item() - 6.1_1_1_2e-1_0      ) < 1e-3
										else:
													# CUDA
													assert abs(result_sum.item() - 4.6_9_3_4_2_8_6_5_0_1_7_0_9_7_2e-0_7      ) < 1e-2
													assert abs(result_mean.item() - 0.0_0_0_2      ) < 1e-3

							def      snake_case							(	self      ):
										if torch_device == "mps":
													return
										__lowerCAmelCase							       =     self.scheduler_classes[0]
										__lowerCAmelCase							       =     self.get_scheduler_config()
										__lowerCAmelCase							       =     scheduler_class(**__a      )

										scheduler.set_timesteps(self.num_inference_steps      )

										__lowerCAmelCase							       =     self.dummy_model()
										__lowerCAmelCase							       =     self.dummy_sample_deter * scheduler.init_noise_sigma
										__lowerCAmelCase							       =     sample.to(__a      )

										for i, t in enumerate(scheduler.timesteps      ):
													__lowerCAmelCase							       =     scheduler.scale_model_input(__a       ,  __a      )

													__lowerCAmelCase							       =     model(__a       ,  __a      )

													__lowerCAmelCase							       =     scheduler.step(__a       ,  __a       ,  __a      )
													__lowerCAmelCase							       =     output.prev_sample

										__lowerCAmelCase							       =     torch.sum(torch.abs(__a      )      )
										__lowerCAmelCase							       =     torch.mean(torch.abs(__a      )      )

										if torch_device in ["cpu", "mps"]:
													assert abs(result_sum.item() - 2_0.4_1_2_5      ) < 1e-2
													assert abs(result_mean.item() - 0.0_2_6_6      ) < 1e-3
										else:
													# CUDA
													assert abs(result_sum.item() - 2_0.4_1_2_5      ) < 1e-2
													assert abs(result_mean.item() - 0.0_2_6_6      ) < 1e-3



							def      snake_case							(	self      ):
										if torch_device == "mps":
													return
										__lowerCAmelCase							       =     self.scheduler_classes[0]
										__lowerCAmelCase							       =     self.get_scheduler_config()
										__lowerCAmelCase							       =     scheduler_class(**__a      )

										scheduler.set_timesteps(self.num_inference_steps       ,  device=__a      )

										__lowerCAmelCase							       =     self.dummy_model()
										__lowerCAmelCase							       =     self.dummy_sample_deter.to(__a      ) * scheduler.init_noise_sigma

										for t in scheduler.timesteps:
													__lowerCAmelCase							       =     scheduler.scale_model_input(__a       ,  __a      )

													__lowerCAmelCase							       =     model(__a       ,  __a      )

													__lowerCAmelCase							       =     scheduler.step(__a       ,  __a       ,  __a      )
													__lowerCAmelCase							       =     output.prev_sample

										__lowerCAmelCase							       =     torch.sum(torch.abs(__a      )      )
										__lowerCAmelCase							       =     torch.mean(torch.abs(__a      )      )

										if str(__a      ).startswith("cpu"      ):
													# The following sum varies between 148 and 156 on mps. Why?
													assert abs(result_sum.item() - 2_0.4_1_2_5      ) < 1e-2
													assert abs(result_mean.item() - 0.0_2_6_6      ) < 1e-3
										else:
													# CUDA
													assert abs(result_sum.item() - 2_0.4_1_2_5      ) < 1e-2
													assert abs(result_mean.item() - 0.0_2_6_6      ) < 1e-3




