





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

      """simple docstring"""


      UpperCamelCase__  		=       ShapEPipeline
      UpperCamelCase__  		=       ["""prompt"""]
      UpperCamelCase__  		=       ["""prompt"""]
      UpperCamelCase__  		=       [
          """num_images_per_prompt""",
          """num_inference_steps""",
          """generator""",
          """latents""",
          """guidance_scale""",
          """frame_size""",
          """output_type""",
          """return_dict""",
      ]
      UpperCamelCase__  		=       False






      @property
      def 			UpperCAmelCase_       (		self:			Any					):
          '''simple docstring'''


          return 32






      @property
      def 			UpperCAmelCase_       (		self:			Any					):
          '''simple docstring'''


          return 32






      @property
      def 			UpperCAmelCase_       (		self:			Optional[Any]					):
          '''simple docstring'''


          return self.time_input_dim * 4






      @property
      def 			UpperCAmelCase_       (		self:			List[Any]					):
          '''simple docstring'''


          return 8






      @property
      def 			UpperCAmelCase_       (		self:			Dict					):
          '''simple docstring'''


          UpperCamelCase__:	List[Any]					=     CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"					)
          return tokenizer






      @property
      def 			UpperCAmelCase_       (		self:			int					):
          '''simple docstring'''


          torch.manual_seed(0					)
          UpperCamelCase__:	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=1000	,    )
          return CLIPTextModelWithProjection(__lowerCamelCase					)






      @property
      def 			UpperCAmelCase_       (		self:			Any					):
          '''simple docstring'''


          torch.manual_seed(0					)

          UpperCamelCase__:	List[str]					=     {
              "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__:	Dict					=     PriorTransformer(**__lowerCamelCase					)
          return model






      @property
      def 			UpperCAmelCase_       (		self:			List[str]					):
          '''simple docstring'''


          torch.manual_seed(0					)

          UpperCamelCase__:	str					=     {
              "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__:	str					=     ShapERenderer(**__lowerCamelCase					)
          return model






      def 			UpperCAmelCase_       (		self:			str					):
          '''simple docstring'''


          UpperCamelCase__:	List[Any]					=     self.dummy_prior
          UpperCamelCase__:	Dict					=     self.dummy_text_encoder
          UpperCamelCase__:	Union[str, Any]					=     self.dummy_tokenizer
          UpperCamelCase__:	Tuple					=     self.dummy_renderer

          UpperCamelCase__:	List[str]					=     HeunDiscreteScheduler(
              beta_schedule="exp"	,    num_train_timesteps=1024	,    prediction_type="sample"	,    use_karras_sigmas=__lowerCamelCase	,    clip_sample=__lowerCamelCase	,    clip_sample_range=1.0	,    )
          UpperCamelCase__:	Optional[int]					=     {
              "prior": prior,
              "text_encoder": text_encoder,
              "tokenizer": tokenizer,
              "renderer": renderer,
              "scheduler": scheduler,
          }

          return components






      def 			UpperCAmelCase_       (		self:			List[str]	,    __lowerCamelCase:			str	,    __lowerCamelCase:			Dict=0					):
          '''simple docstring'''


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






      def 			UpperCAmelCase_       (		self:			Optional[Any]					):
          '''simple docstring'''


          UpperCamelCase__:	str					=     "cpu"

          UpperCamelCase__:	Dict					=     self.get_dummy_components()

          UpperCamelCase__:	Any					=     self.pipeline_class(**__lowerCamelCase					)
          UpperCamelCase__:	List[str]					=     pipe.to(__lowerCamelCase					)

          pipe.set_progress_bar_config(disable=__lowerCamelCase					)

          UpperCamelCase__:	List[Any]					=     pipe(**self.get_dummy_inputs(__lowerCamelCase					)					)
          UpperCamelCase__:	Any					=     output.images[0]
          UpperCamelCase__:	Tuple					=     image[0, -3:, -3:, -1]

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

          UpperCamelCase__:	Any					=     np.array(
              [
                  0.00_039_216,
                  0.00_039_216,
                  0.00_039_216,
                  0.00_039_216,
                  0.00_039_216,
                  0.00_039_216,
                  0.00_039_216,
                  0.00_039_216,
                  0.00_039_216,
              ]					)

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






      def 			UpperCAmelCase_       (		self:			Optional[int]					):
          '''simple docstring'''


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






      def 			UpperCAmelCase_       (		self:			Union[str, Any]					):
          '''simple docstring'''


          UpperCamelCase__:	Tuple					=     torch_device == "cpu"
          UpperCamelCase__:	Tuple					=     True

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






      def 			UpperCAmelCase_       (		self:			Optional[Any]					):
          '''simple docstring'''


          UpperCamelCase__:	str					=     self.get_dummy_components()
          UpperCamelCase__:	int					=     self.pipeline_class(**__lowerCamelCase					)
          UpperCamelCase__:	Any					=     pipe.to(__lowerCamelCase					)
          pipe.set_progress_bar_config(disable=__lowerCamelCase					)

          UpperCamelCase__:	str					=     1
          UpperCamelCase__:	str					=     2

          UpperCamelCase__:	int					=     self.get_dummy_inputs(__lowerCamelCase					)

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

          UpperCamelCase__:	Any					=     pipe(**__lowerCamelCase	,    num_images_per_prompt=__lowerCamelCase					)[0]

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




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

      """simple docstring"""


      def 			UpperCAmelCase_       (		self:			Dict					):
          '''simple docstring'''


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






      def 			UpperCAmelCase_       (		self:			str					):
          '''simple docstring'''


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

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

          UpperCamelCase__:	List[Any]					=     pipe(
              "a shark"	,    generator=__lowerCamelCase	,    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(__lowerCamelCase	,    __lowerCamelCase					)








import unittest

import numpy as np
from datasets import load_dataset

from transformers.testing_utils import require_torch, require_vision
from transformers.utils import is_torch_available, is_vision_available

from ...test_image_processing_common import ImageProcessingSavingTestMixin, prepare_image_inputs


if is_torch_available():
   import torch

if is_vision_available():
   from PIL import Image

   from transformers import BeitImageProcessor




class        _a    (    unittest.TestCase):

      """simple docstring"""


      def __init__(		self:			Dict	,    __lowerCamelCase:			Any	,    __lowerCamelCase:			Optional[int]=7	,    __lowerCamelCase:			Any=3	,    __lowerCamelCase:			List[str]=18	,    __lowerCamelCase:			List[Any]=30	,    __lowerCamelCase:			Tuple=400	,    __lowerCamelCase:			List[str]=True	,    __lowerCamelCase:			Any=None	,    __lowerCamelCase:			int=True	,    __lowerCamelCase:			Any=None	,    __lowerCamelCase:			Dict=True	,    __lowerCamelCase:			List[Any]=[0.5, 0.5, 0.5]	,    __lowerCamelCase:			Optional[Any]=[0.5, 0.5, 0.5]	,    __lowerCamelCase:			int=False	,    ):
          '''simple docstring'''


          UpperCamelCase__:	Optional[int]					=     size if size is not None else {"height": 20, "width": 20}
          UpperCamelCase__:	List[Any]					=     crop_size if crop_size is not None else {"height": 18, "width": 18}
          UpperCamelCase__:	Optional[int]					=     parent
          UpperCamelCase__:	int					=     batch_size
          UpperCamelCase__:	int					=     num_channels
          UpperCamelCase__:	str					=     image_size
          UpperCamelCase__:	Any					=     min_resolution
          UpperCamelCase__:	Union[str, Any]					=     max_resolution
          UpperCamelCase__:	Optional[Any]					=     do_resize
          UpperCamelCase__:	Any					=     size
          UpperCamelCase__:	str					=     do_center_crop
          UpperCamelCase__:	Any					=     crop_size
          UpperCamelCase__:	Any					=     do_normalize
          UpperCamelCase__:	int					=     image_mean
          UpperCamelCase__:	Tuple					=     image_std
          UpperCamelCase__:	int					=     do_reduce_labels






      def 			UpperCAmelCase_       (		self:			Tuple					):
          '''simple docstring'''


          return {
              "do_resize": self.do_resize,
              "size": self.size,
              "do_center_crop": self.do_center_crop,
              "crop_size": self.crop_size,
              "do_normalize": self.do_normalize,
              "image_mean": self.image_mean,
              "image_std": self.image_std,
              "do_reduce_labels": self.do_reduce_labels,
          }



def 				lowerCAmelCase_ (							):
    UpperCamelCase__:	Dict					=     load_dataset("hf-internal-testing/fixtures_ade20k"    ,split="test")

    UpperCamelCase__:	Optional[Any]					=     Image.open(dataset[0]["file"])
    UpperCamelCase__:	str					=     Image.open(dataset[1]["file"])

    return image, map



def 				lowerCAmelCase_ (							):
    UpperCamelCase__:	Dict					=     load_dataset("hf-internal-testing/fixtures_ade20k"    ,split="test")

    UpperCamelCase__:	int					=     Image.open(ds[0]["file"])
    UpperCamelCase__:	int					=     Image.open(ds[1]["file"])
    UpperCamelCase__:	List[str]					=     Image.open(ds[2]["file"])
    UpperCamelCase__:	List[Any]					=     Image.open(ds[3]["file"])

    return [imagea, imagea], [mapa, mapa]




@require_torch
@require_vision
class        _a    (    UpperCamelCase__						,			unittest.TestCase):

      """simple docstring"""


      UpperCamelCase__  		=       BeitImageProcessor if is_vision_available() else None






      def 			UpperCAmelCase_       (		self:			Optional[int]					):
          '''simple docstring'''


          UpperCamelCase__:	str					=     BeitImageProcessingTester(self					)






      @property
      def 			UpperCAmelCase_       (		self:			Any					):
          '''simple docstring'''


          return self.image_processor_tester.prepare_image_processor_dict()






      def 			UpperCAmelCase_       (		self:			Dict					):
          '''simple docstring'''


          UpperCamelCase__:	int					=     self.image_processing_class(**self.image_processor_dict					)
          self.assertTrue(hasattr(__lowerCamelCase	,    "do_resize"					)					)
          self.assertTrue(hasattr(__lowerCamelCase	,    "size"					)					)
          self.assertTrue(hasattr(__lowerCamelCase	,    "do_center_crop"					)					)
          self.assertTrue(hasattr(__lowerCamelCase	,    "center_crop"					)					)
          self.assertTrue(hasattr(__lowerCamelCase	,    "do_normalize"					)					)
          self.assertTrue(hasattr(__lowerCamelCase	,    "image_mean"					)					)
          self.assertTrue(hasattr(__lowerCamelCase	,    "image_std"					)					)






      def 			UpperCAmelCase_       (		self:			Optional[int]					):
          '''simple docstring'''


          UpperCamelCase__:	Tuple					=     self.image_processing_class.from_dict(self.image_processor_dict					)
          self.assertEqual(image_processor.size	,    {"height": 20, "width": 20}					)
          self.assertEqual(image_processor.crop_size	,    {"height": 18, "width": 18}					)
          self.assertEqual(image_processor.do_reduce_labels	,    __lowerCamelCase					)

          UpperCamelCase__:	int					=     self.image_processing_class.from_dict(
              self.image_processor_dict	,    size=42	,    crop_size=84	,    reduce_labels=__lowerCamelCase					)
          self.assertEqual(image_processor.size	,    {"height": 42, "width": 42}					)
          self.assertEqual(image_processor.crop_size	,    {"height": 84, "width": 84}					)
          self.assertEqual(image_processor.do_reduce_labels	,    __lowerCamelCase					)






      def 			UpperCAmelCase_       (		self:			Tuple					):
          '''simple docstring'''


          pass






      def 			UpperCAmelCase_       (		self:			Dict					):
          '''simple docstring'''


          UpperCamelCase__:	Tuple					=     self.image_processing_class(**self.image_processor_dict					)
          # create random PIL images
          UpperCamelCase__:	Optional[Any]					=     prepare_image_inputs(self.image_processor_tester	,    equal_resolution=__lowerCamelCase					)
          for image in image_inputs:
              self.assertIsInstance(__lowerCamelCase	,    Image.Image					)

          # Test not batched input
          UpperCamelCase__:	List[Any]					=     image_processing(image_inputs[0]	,    return_tensors="pt"					).pixel_values
          self.assertEqual(
              encoded_images.shape	,    (
                  1,
                  self.image_processor_tester.num_channels,
                  self.image_processor_tester.crop_size["height"],
                  self.image_processor_tester.crop_size["width"],
              )	,    )

          # Test batched
          UpperCamelCase__:	str					=     image_processing(__lowerCamelCase	,    return_tensors="pt"					).pixel_values
          self.assertEqual(
              encoded_images.shape	,    (
                  self.image_processor_tester.batch_size,
                  self.image_processor_tester.num_channels,
                  self.image_processor_tester.crop_size["height"],
                  self.image_processor_tester.crop_size["width"],
              )	,    )






      def 			UpperCAmelCase_       (		self:			Dict					):
          '''simple docstring'''


          UpperCamelCase__:	Optional[Any]					=     self.image_processing_class(**self.image_processor_dict					)
          # create random numpy tensors
          UpperCamelCase__:	List[str]					=     prepare_image_inputs(self.image_processor_tester	,    equal_resolution=__lowerCamelCase	,    numpify=__lowerCamelCase					)
          for image in image_inputs:
              self.assertIsInstance(__lowerCamelCase	,    np.ndarray					)

          # Test not batched input
          UpperCamelCase__:	Union[str, Any]					=     image_processing(image_inputs[0]	,    return_tensors="pt"					).pixel_values
          self.assertEqual(
              encoded_images.shape	,    (
                  1,
                  self.image_processor_tester.num_channels,
                  self.image_processor_tester.crop_size["height"],
                  self.image_processor_tester.crop_size["width"],
              )	,    )

          # Test batched
          UpperCamelCase__:	Dict					=     image_processing(__lowerCamelCase	,    return_tensors="pt"					).pixel_values
          self.assertEqual(
              encoded_images.shape	,    (
                  self.image_processor_tester.batch_size,
                  self.image_processor_tester.num_channels,
                  self.image_processor_tester.crop_size["height"],
                  self.image_processor_tester.crop_size["width"],
              )	,    )






      def 			UpperCAmelCase_       (		self:			Union[str, Any]					):
          '''simple docstring'''


          UpperCamelCase__:	int					=     self.image_processing_class(**self.image_processor_dict					)
          # create random PyTorch tensors
          UpperCamelCase__:	int					=     prepare_image_inputs(self.image_processor_tester	,    equal_resolution=__lowerCamelCase	,    torchify=__lowerCamelCase					)
          for image in image_inputs:
              self.assertIsInstance(__lowerCamelCase	,    torch.Tensor					)

          # Test not batched input
          UpperCamelCase__:	Dict					=     image_processing(image_inputs[0]	,    return_tensors="pt"					).pixel_values
          self.assertEqual(
              encoded_images.shape	,    (
                  1,
                  self.image_processor_tester.num_channels,
                  self.image_processor_tester.crop_size["height"],
                  self.image_processor_tester.crop_size["width"],
              )	,    )

          # Test batched
          UpperCamelCase__:	str					=     image_processing(__lowerCamelCase	,    return_tensors="pt"					).pixel_values
          self.assertEqual(
              encoded_images.shape	,    (
                  self.image_processor_tester.batch_size,
                  self.image_processor_tester.num_channels,
                  self.image_processor_tester.crop_size["height"],
                  self.image_processor_tester.crop_size["width"],
              )	,    )






      def 			UpperCAmelCase_       (		self:			Tuple					):
          '''simple docstring'''


          UpperCamelCase__:	Union[str, Any]					=     self.image_processing_class(**self.image_processor_dict					)
          # create random PyTorch tensors
          UpperCamelCase__:	Dict					=     prepare_image_inputs(self.image_processor_tester	,    equal_resolution=__lowerCamelCase	,    torchify=__lowerCamelCase					)
          UpperCamelCase__:	str					=     []
          for image in image_inputs:
              self.assertIsInstance(__lowerCamelCase	,    torch.Tensor					)
              maps.append(torch.zeros(image.shape[-2:]					).long()					)

          # Test not batched input
          UpperCamelCase__:	Dict					=     image_processing(image_inputs[0]	,    maps[0]	,    return_tensors="pt"					)
          self.assertEqual(
              encoding["pixel_values"].shape	,    (
                  1,
                  self.image_processor_tester.num_channels,
                  self.image_processor_tester.crop_size["height"],
                  self.image_processor_tester.crop_size["width"],
              )	,    )
          self.assertEqual(
              encoding["labels"].shape	,    (
                  1,
                  self.image_processor_tester.crop_size["height"],
                  self.image_processor_tester.crop_size["width"],
              )	,    )
          self.assertEqual(encoding["labels"].dtype	,    torch.long					)
          self.assertTrue(encoding["labels"].min().item() >= 0					)
          self.assertTrue(encoding["labels"].max().item() <= 255					)

          # Test batched
          UpperCamelCase__:	Any					=     image_processing(__lowerCamelCase	,    __lowerCamelCase	,    return_tensors="pt"					)
          self.assertEqual(
              encoding["pixel_values"].shape	,    (
                  self.image_processor_tester.batch_size,
                  self.image_processor_tester.num_channels,
                  self.image_processor_tester.crop_size["height"],
                  self.image_processor_tester.crop_size["width"],
              )	,    )
          self.assertEqual(
              encoding["labels"].shape	,    (
                  self.image_processor_tester.batch_size,
                  self.image_processor_tester.crop_size["height"],
                  self.image_processor_tester.crop_size["width"],
              )	,    )
          self.assertEqual(encoding["labels"].dtype	,    torch.long					)
          self.assertTrue(encoding["labels"].min().item() >= 0					)
          self.assertTrue(encoding["labels"].max().item() <= 255					)

          # Test not batched input (PIL images)
          UpperCamelCase__				, UpperCamelCase__:	str					=     prepare_semantic_single_inputs()

          UpperCamelCase__:	Any					=     image_processing(__lowerCamelCase	,    __lowerCamelCase	,    return_tensors="pt"					)
          self.assertEqual(
              encoding["pixel_values"].shape	,    (
                  1,
                  self.image_processor_tester.num_channels,
                  self.image_processor_tester.crop_size["height"],
                  self.image_processor_tester.crop_size["width"],
              )	,    )
          self.assertEqual(
              encoding["labels"].shape	,    (
                  1,
                  self.image_processor_tester.crop_size["height"],
                  self.image_processor_tester.crop_size["width"],
              )	,    )
          self.assertEqual(encoding["labels"].dtype	,    torch.long					)
          self.assertTrue(encoding["labels"].min().item() >= 0					)
          self.assertTrue(encoding["labels"].max().item() <= 255					)

          # Test batched input (PIL images)
          UpperCamelCase__				, UpperCamelCase__:	List[str]					=     prepare_semantic_batch_inputs()

          UpperCamelCase__:	Optional[int]					=     image_processing(__lowerCamelCase	,    __lowerCamelCase	,    return_tensors="pt"					)
          self.assertEqual(
              encoding["pixel_values"].shape	,    (
                  2,
                  self.image_processor_tester.num_channels,
                  self.image_processor_tester.crop_size["height"],
                  self.image_processor_tester.crop_size["width"],
              )	,    )
          self.assertEqual(
              encoding["labels"].shape	,    (
                  2,
                  self.image_processor_tester.crop_size["height"],
                  self.image_processor_tester.crop_size["width"],
              )	,    )
          self.assertEqual(encoding["labels"].dtype	,    torch.long					)
          self.assertTrue(encoding["labels"].min().item() >= 0					)
          self.assertTrue(encoding["labels"].max().item() <= 255					)






      def 			UpperCAmelCase_       (		self:			Any					):
          '''simple docstring'''


          UpperCamelCase__:	Dict					=     self.image_processing_class(**self.image_processor_dict					)

          # ADE20k has 150 classes, and the background is included, so labels should be between 0 and 150
          UpperCamelCase__				, UpperCamelCase__:	Any					=     prepare_semantic_single_inputs()
          UpperCamelCase__:	int					=     image_processing(__lowerCamelCase	,    __lowerCamelCase	,    return_tensors="pt"					)
          self.assertTrue(encoding["labels"].min().item() >= 0					)
          self.assertTrue(encoding["labels"].max().item() <= 150					)

          UpperCamelCase__:	List[Any]					=     True
          UpperCamelCase__:	List[str]					=     image_processing(__lowerCamelCase	,    __lowerCamelCase	,    return_tensors="pt"					)
          self.assertTrue(encoding["labels"].min().item() >= 0					)
          self.assertTrue(encoding["labels"].max().item() <= 255					)

