





'''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	(       a_			,       unittest.TestCase				):
 __lowercase   		=				ShapEPipeline
 __lowercase   		=				["""prompt"""]
 __lowercase   		=				["""prompt"""]
 __lowercase   		=				[
     """num_images_per_prompt""",
     """num_inference_steps""",
     """generator""",
     """latents""",
     """guidance_scale""",
     """frame_size""",
     """output_type""",
     """return_dict""",
 ]
 __lowercase   		=				False
 @property
 def      UpperCAmelCase_							(				self      :int							)->							Any:
     return 32
 @property
 def      UpperCAmelCase_							(				self      :List[str]							)->							Dict:
     return 32
 @property
 def      UpperCAmelCase_							(				self      :List[str]							)->							Any:
     return self.time_input_dim * 4
 @property
 def      UpperCAmelCase_							(				self      :Union[str, Any]							)->							Optional[int]:
     return 8
 @property
 def      UpperCAmelCase_							(				self      :List[Any]							)->							Any:
     A__				=			CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"							)
     return tokenizer
 @property
 def      UpperCAmelCase_							(				self      :Optional[int]							)->							Optional[int]:
     torch.manual_seed(0							)
     A__				=			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      UpperCAmelCase_							(				self      :List[str]							)->							List[str]:
     torch.manual_seed(0							)

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

     A__				=			PriorTransformer(**_a							)
     return model
 @property
 def      UpperCAmelCase_							(				self      :Union[str, Any]							)->							str:
     torch.manual_seed(0							)

     A__				=			{
         "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,
         ),
     }
     A__				=			ShapERenderer(**_a							)
     return model
 def      UpperCAmelCase_							(				self      :int							)->							List[str]:
     A__				=			self.dummy_prior
     A__				=			self.dummy_text_encoder
     A__				=			self.dummy_tokenizer
     A__				=			self.dummy_renderer

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

     return components
 def      UpperCAmelCase_							(				self      :Optional[int]						,       lowercase_      :int						,       lowercase_      :Optional[int]=0							)->							List[Any]:
     if str(_a							).startswith("mps"							):
         A__				=			torch.manual_seed(_a							)
     else:
         A__				=			torch.Generator(device=_a							).manual_seed(_a							)
     A__				=			{
         "prompt": "horse",
         "generator": generator,
         "num_inference_steps": 1,
         "frame_size": 32,
         "output_type": "np",
     }
     return inputs
 def      UpperCAmelCase_							(				self      :List[Any]							)->							int:
     A__				=			"cpu"

     A__				=			self.get_dummy_components()

     A__				=			self.pipeline_class(**_a							)
     A__				=			pipe.to(_a							)

     pipe.set_progress_bar_config(disable=_a							)

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

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

     A__				=			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      UpperCAmelCase_							(				self      :Any							)->							Dict:
     # NOTE: Larger batch sizes cause this test to timeout, only test on smaller batches
     self._test_inference_batch_consistent(batch_sizes=[1, 2]							)
 def      UpperCAmelCase_							(				self      :List[str]							)->							Union[str, Any]:
     A__				=			torch_device == "cpu"
     A__				=			True

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



 def      UpperCAmelCase_							(				self      :Optional[Any]							)->							List[str]:
     A__				=			self.get_dummy_components()
     A__				=			self.pipeline_class(**_a							)
     A__				=			pipe.to(_a							)
     pipe.set_progress_bar_config(disable=_a							)

     A__				=			1
     A__				=			2

     A__				=			self.get_dummy_inputs(_a							)

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

     A__				=			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      UpperCAmelCase_							(				self      :Union[str, Any]							)->							str:
     # clean up the VRAM after each test
     super().tearDown()
     gc.collect()
     torch.cuda.empty_cache()



 def      UpperCAmelCase_							(				self      :List[str]							)->							Optional[int]:
     A__				=			load_numpy(
         "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
         "/shap_e/test_shap_e_np_out.npy"							)
     A__				=			ShapEPipeline.from_pretrained("openai/shap-e"							)
     A__				=			pipe.to(_a							)
     pipe.set_progress_bar_config(disable=_a							)

     A__				=			torch.Generator(device=_a							).manual_seed(0							)

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





"""simple docstring"""




from typing import Dict, List, Optional, Tuple, Union

import numpy as np

from ...image_processing_utils import BaseImageProcessor, BatchFeature, get_size_dict
from ...image_transforms import (
    center_crop,
    flip_channel_order,
    get_resize_output_image_size,
    rescale,
    resize,
    to_channel_dimension_format,
)
from ...image_utils import (
    ChannelDimension,
    ImageInput,
    PILImageResampling,
    make_list_of_images,
    to_numpy_array,
    valid_images,
)
from ...utils import TensorType, is_torch_available, is_torch_tensor, is_vision_available, logging


if is_vision_available():
   import PIL

if is_torch_available():
   import torch


_A			:       List[Any]   		=       logging.get_logger(__name__)

class 		a__						(      a_						):
   __lowerCAmelCase            =				["""pixel_values"""]
   def __init__(			self	,    _a = True	,    _a = None	,    _a = PILImageResampling.BILINEAR	,    _a = True	,    _a = 1 / 255	,    _a = True	,    _a = None	,    _a = True	,    **_a	,    ):
          super().__init__(**_a				)
          lowercase    :				Optional[Any]       		=							size if size is not None else {"shortest_edge": 224}
          lowercase    :				List[Any]       		=							get_size_dict(_a	,    default_to_square=_a				)
          lowercase    :				str       		=							crop_size if crop_size is not None else {"height": 256, "width": 256}
          lowercase    :				List[str]       		=							get_size_dict(_a	,    param_name="crop_size"				)

          lowercase    :				int       		=							do_resize
          lowercase    :				Optional[int]       		=							size
          lowercase    :				str       		=							resample
          lowercase    :				List[Any]       		=							do_rescale
          lowercase    :				Union[str, Any]       		=							rescale_factor
          lowercase    :				Optional[int]       		=							do_center_crop
          lowercase    :				Union[str, Any]       		=							crop_size
          lowercase    :				Optional[Any]       		=							do_flip_channel_order
   def 			__magic_name__			(			self	,    _a	,    _a	,    _a = PIL.Image.BILINEAR	,    _a = None	,    **_a	,    ):
          lowercase    :				List[Any]       		=							get_size_dict(_a	,    default_to_square=_a				)
          if "shortest_edge" not in size:
                 raise ValueError(f"""The `size` dictionary must contain the key `shortest_edge`. Got {size.keys()}"""				)
          lowercase    :				Union[str, Any]       		=							get_resize_output_image_size(_a	,    size=size["shortest_edge"]	,    default_to_square=_a				)
          return resize(_a	,    size=_a	,    resample=_a	,    data_format=_a	,    **_a				)
   def 			__magic_name__			(			self	,    _a	,    _a	,    _a = None	,    **_a	,    ):
          lowercase    :				str       		=							get_size_dict(_a				)
          if "height" not in size or "width" not in size:
                 raise ValueError(f"""The `size` dictionary must contain the keys `height` and `width`. Got {size.keys()}"""				)
          return center_crop(_a	,    size=(size["height"], size["width"])	,    data_format=_a	,    **_a				)
   def 			__magic_name__			(			self	,    _a	,    _a	,    _a = None	,    **_a	,    ):
          return rescale(_a	,    scale=_a	,    data_format=_a	,    **_a				)
   def 			__magic_name__			(			self	,    _a	,    _a = None				):
          return flip_channel_order(_a	,    data_format=_a				)
   def 			__magic_name__			(			self	,    _a	,    _a = None	,    _a = None	,    _a = None	,    _a = None	,    _a = None	,    _a = None	,    _a = None	,    _a = None	,    _a = None	,    _a = ChannelDimension.FIRST	,    **_a	,    ):
          lowercase    :				Optional[Any]       		=							do_resize if do_resize is not None else self.do_resize
          lowercase    :				Tuple       		=							resample if resample is not None else self.resample
          lowercase    :				List[str]       		=							do_rescale if do_rescale is not None else self.do_rescale
          lowercase    :				List[Any]       		=							rescale_factor if rescale_factor is not None else self.rescale_factor
          lowercase    :				Union[str, Any]       		=							do_center_crop if do_center_crop is not None else self.do_center_crop
          lowercase    :				Optional[int]       		=							(
              do_flip_channel_order if do_flip_channel_order is not None else self.do_flip_channel_order
          )

          lowercase    :				str       		=							size if size is not None else self.size
          lowercase    :				Any       		=							get_size_dict(_a	,    default_to_square=_a				)
          lowercase    :				int       		=							crop_size if crop_size is not None else self.crop_size
          lowercase    :				Any       		=							get_size_dict(_a	,    param_name="crop_size"				)

          lowercase    :				int       		=							make_list_of_images(_a				)

          if not valid_images(_a				):
                 raise ValueError(
                     "Invalid image type. Must be of type PIL.Image.Image, numpy.ndarray, "
                     "torch.Tensor, tf.Tensor or jax.ndarray."				)

          if do_resize and size is None:
                 raise ValueError("Size must be specified if do_resize is True."				)

          if do_rescale and rescale_factor is None:
                 raise ValueError("Rescale factor must be specified if do_rescale is True."				)

          if do_center_crop and crop_size is None:
                 raise ValueError("Crop size must be specified if do_center_crop is True."				)

          # All transformations expect numpy arrays.
          lowercase    :				Any       		=							[to_numpy_array(_a				) for image in images]

          if do_resize:
                 lowercase    :				Optional[int]       		=							[self.resize(image=_a	,    size=_a	,    resample=_a				) for image in images]

          if do_center_crop:
                 lowercase    :				str       		=							[self.center_crop(image=_a	,    size=_a				) for image in images]

          if do_rescale:
                 lowercase    :				Union[str, Any]       		=							[self.rescale(image=_a	,    scale=_a				) for image in images]

          # the pretrained checkpoints assume images are BGR, not RGB
          if do_flip_channel_order:
                 lowercase    :				int       		=							[self.flip_channel_order(image=_a				) for image in images]

          lowercase    :				int       		=							[to_channel_dimension_format(_a	,    _a				) for image in images]

          lowercase    :				Optional[Any]       		=							{"pixel_values": images}
          return BatchFeature(data=_a	,    tensor_type=_a				)

   def 			__magic_name__			(			self	,    _a	,    _a = None				):
          lowercase    :				Optional[int]       		=							outputs.logits

          # Resize logits and compute semantic segmentation maps
          if target_sizes is not None:
                 if len(_a				) != len(_a				):
                        raise ValueError(
                            "Make sure that you pass in as many target sizes as the batch dimension of the logits"				)

                 if is_torch_tensor(_a				):
                        lowercase    :				Tuple       		=							target_sizes.numpy()

                 lowercase    :				List[Any]       		=							[]

                 for idx in range(len(_a				)				):
                        lowercase    :				Dict       		=							torch.nn.functional.interpolate(
                            logits[idx].unsqueeze(dim=0				)	,    size=target_sizes[idx]	,    mode="bilinear"	,    align_corners=_a				)
                        lowercase    :				Tuple       		=							resized_logits[0].argmax(dim=0				)
                        semantic_segmentation.append(_a				)
          else:
                 lowercase    :				str       		=							logits.argmax(dim=1				)
                 lowercase    :				Dict       		=							[semantic_segmentation[i] for i in range(semantic_segmentation.shape[0]				)]

          return semantic_segmentation


