'''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 		__A       (__magic_name__ ,					unittest.TestCase   ):
       snake_case       :Any				    =					ShapEPipeline
       snake_case       :Tuple				    =					["prompt"]
       snake_case       :Dict				    =					["prompt"]
       snake_case       :Tuple				    =					[
           "num_images_per_prompt",
           "num_inference_steps",
           "generator",
           "latents",
           "guidance_scale",
           "frame_size",
           "output_type",
           "return_dict",
       ]
       snake_case       :Any				    =					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      ):
              __UpperCAmelCase  :				Optional[Any]       				=		CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"      )
              return tokenizer






       @property
       def 			_snake_case						(      self      ):
              torch.manual_seed(0      )
              __UpperCAmelCase  :				List[str]       				=		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(UpperCamelCase_      )






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

              __UpperCAmelCase  :				Optional[Any]       				=		{
                  "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,
              }

              __UpperCAmelCase  :				int       				=		PriorTransformer(**UpperCamelCase_      )
              return model






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

              __UpperCAmelCase  :				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,
                  ),
              }
              __UpperCAmelCase  :				Optional[Any]       				=		ShapERenderer(**UpperCamelCase_      )
              return model






       def 			_snake_case						(      self      ):
              __UpperCAmelCase  :				Any       				=		self.dummy_prior
              __UpperCAmelCase  :				Tuple       				=		self.dummy_text_encoder
              __UpperCAmelCase  :				List[Any]       				=		self.dummy_tokenizer
              __UpperCAmelCase  :				Dict       				=		self.dummy_renderer

              __UpperCAmelCase  :				Optional[Any]       				=		HeunDiscreteScheduler(
                  beta_schedule="exp"				,     num_train_timesteps=10_24				,     prediction_type="sample"				,     use_karras_sigmas=UpperCamelCase_				,     clip_sample=UpperCamelCase_				,     clip_sample_range=1.0				,     )
              __UpperCAmelCase  :				Optional[int]       				=		{
                  "prior": prior,
                  "text_encoder": text_encoder,
                  "tokenizer": tokenizer,
                  "renderer": renderer,
                  "scheduler": scheduler,
              }

              return components






       def 			_snake_case						(      self				,     UpperCamelCase_				,     UpperCamelCase_=0      ):
              if str(UpperCamelCase_      ).startswith("mps"      ):
                     __UpperCAmelCase  :				Optional[Any]       				=		torch.manual_seed(UpperCamelCase_      )
              else:
                     __UpperCAmelCase  :				Any       				=		torch.Generator(device=UpperCamelCase_      ).manual_seed(UpperCamelCase_      )
              __UpperCAmelCase  :				List[str]       				=		{
                  "prompt": "horse",
                  "generator": generator,
                  "num_inference_steps": 1,
                  "frame_size": 32,
                  "output_type": "np",
              }
              return inputs






       def 			_snake_case						(      self      ):
              __UpperCAmelCase  :				Optional[Any]       				=		"cpu"

              __UpperCAmelCase  :				List[str]       				=		self.get_dummy_components()

              __UpperCAmelCase  :				Optional[int]       				=		self.pipeline_class(**UpperCamelCase_      )
              __UpperCAmelCase  :				Tuple       				=		pipe.to(UpperCamelCase_      )

              pipe.set_progress_bar_config(disable=UpperCamelCase_      )

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

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

              __UpperCAmelCase  :				Union[str, Any]       				=		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      ):
              __UpperCAmelCase  :				int       				=		torch_device == "cpu"
              __UpperCAmelCase  :				Optional[Any]       				=		True

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







       def 			_snake_case						(      self      ):
              __UpperCAmelCase  :				List[str]       				=		self.get_dummy_components()
              __UpperCAmelCase  :				Optional[Any]       				=		self.pipeline_class(**UpperCamelCase_      )
              __UpperCAmelCase  :				int       				=		pipe.to(UpperCamelCase_      )
              pipe.set_progress_bar_config(disable=UpperCamelCase_      )

              __UpperCAmelCase  :				Dict       				=		1
              __UpperCAmelCase  :				Optional[Any]       				=		2

              __UpperCAmelCase  :				List[Any]       				=		self.get_dummy_inputs(UpperCamelCase_      )

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

              __UpperCAmelCase  :				Any       				=		pipe(**UpperCamelCase_				,     num_images_per_prompt=UpperCamelCase_      )[0]

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






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






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







       def 			_snake_case						(      self      ):
              __UpperCAmelCase  :				Optional[Any]       				=		load_numpy(
                  "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
                  "/shap_e/test_shap_e_np_out.npy"      )
              __UpperCAmelCase  :				Tuple       				=		ShapEPipeline.from_pretrained("openai/shap-e"      )
              __UpperCAmelCase  :				Dict       				=		pipe.to(UpperCamelCase_      )
              pipe.set_progress_bar_config(disable=UpperCamelCase_      )

              __UpperCAmelCase  :				Union[str, Any]       				=		torch.Generator(device=UpperCamelCase_      ).manual_seed(0      )

              __UpperCAmelCase  :				Union[str, Any]       				=		pipe(
                  "a shark"				,     generator=UpperCamelCase_				,     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(UpperCamelCase_				,     UpperCamelCase_      )







'''simple docstring'''




import time
import unittest

from transformers import is_torch_available
from transformers.testing_utils import require_torch, torch_device

from ..test_modeling_common import ids_tensor


if is_torch_available():
   import torch

   from transformers.generation import (
       MaxLengthCriteria,
       MaxNewTokensCriteria,
       MaxTimeCriteria,
       StoppingCriteriaList,
       validate_stopping_criteria,
   )






@require_torch
class 		__A       (unittest.TestCase   ):






       def 			_snake_case						(      self				,     UpperCamelCase_      ):
              __UpperCAmelCase  :				List[str]       				=		3
              __UpperCAmelCase  :				Tuple       				=		2_50

              __UpperCAmelCase  :				str       				=		ids_tensor((batch_size, length)				,     UpperCamelCase_      )
              __UpperCAmelCase  :				Any       				=		torch.ones((batch_size, length)				,     device=UpperCamelCase_				,     dtype=torch.float      ) / length
              return input_ids, scores






       def 			_snake_case						(      self      ):
              __UpperCAmelCase				, __UpperCAmelCase  :				Tuple       				=		self._get_tensors(5      )

              __UpperCAmelCase  :				Tuple       				=		StoppingCriteriaList(
                  [
                      MaxLengthCriteria(max_length=10      ),
                      MaxTimeCriteria(max_time=0.1      ),
                  ]      )

              self.assertFalse(criteria(UpperCamelCase_				,     UpperCamelCase_      )      )

              __UpperCAmelCase				, __UpperCAmelCase  :				int       				=		self._get_tensors(9      )
              self.assertFalse(criteria(UpperCamelCase_				,     UpperCamelCase_      )      )

              __UpperCAmelCase				, __UpperCAmelCase  :				Optional[int]       				=		self._get_tensors(10      )
              self.assertTrue(criteria(UpperCamelCase_				,     UpperCamelCase_      )      )






       def 			_snake_case						(      self      ):
              __UpperCAmelCase  :				int       				=		MaxLengthCriteria(max_length=10      )

              __UpperCAmelCase				, __UpperCAmelCase  :				Tuple       				=		self._get_tensors(5      )
              self.assertFalse(criteria(UpperCamelCase_				,     UpperCamelCase_      )      )

              __UpperCAmelCase				, __UpperCAmelCase  :				Dict       				=		self._get_tensors(9      )
              self.assertFalse(criteria(UpperCamelCase_				,     UpperCamelCase_      )      )

              __UpperCAmelCase				, __UpperCAmelCase  :				Optional[int]       				=		self._get_tensors(10      )
              self.assertTrue(criteria(UpperCamelCase_				,     UpperCamelCase_      )      )






       def 			_snake_case						(      self      ):
              __UpperCAmelCase  :				Optional[Any]       				=		MaxNewTokensCriteria(start_length=5				,     max_new_tokens=5      )

              __UpperCAmelCase				, __UpperCAmelCase  :				List[str]       				=		self._get_tensors(5      )
              self.assertFalse(criteria(UpperCamelCase_				,     UpperCamelCase_      )      )

              __UpperCAmelCase				, __UpperCAmelCase  :				Dict       				=		self._get_tensors(9      )
              self.assertFalse(criteria(UpperCamelCase_				,     UpperCamelCase_      )      )

              __UpperCAmelCase				, __UpperCAmelCase  :				Optional[Any]       				=		self._get_tensors(10      )
              self.assertTrue(criteria(UpperCamelCase_				,     UpperCamelCase_      )      )

              __UpperCAmelCase  :				Union[str, Any]       				=		StoppingCriteriaList([criteria]      )
              self.assertEqual(criteria_list.max_length				,     10      )






       def 			_snake_case						(      self      ):
              __UpperCAmelCase				, __UpperCAmelCase  :				Optional[Any]       				=		self._get_tensors(5      )

              __UpperCAmelCase  :				str       				=		MaxTimeCriteria(max_time=0.1      )
              self.assertFalse(criteria(UpperCamelCase_				,     UpperCamelCase_      )      )

              __UpperCAmelCase  :				str       				=		MaxTimeCriteria(max_time=0.1				,     initial_timestamp=time.time() - 0.2      )
              self.assertTrue(criteria(UpperCamelCase_				,     UpperCamelCase_      )      )







       def 			_snake_case						(      self      ):
              validate_stopping_criteria(StoppingCriteriaList([MaxLengthCriteria(10      )]      )				,     10      )

              with self.assertWarns(UpperCamelCase_      ):
                     validate_stopping_criteria(StoppingCriteriaList([MaxLengthCriteria(10      )]      )				,     11      )

              __UpperCAmelCase  :				Optional[int]       				=		validate_stopping_criteria(StoppingCriteriaList()				,     11      )

              self.assertEqual(len(UpperCamelCase_      )				,     1      )






