
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 					lowerCamelCase		(  __lowercase							,			unittest.TestCase  ):
     UpperCAmelCase__    :							Any									=		ShapEPipeline
     UpperCAmelCase__    :							Any									=		["prompt"]
     UpperCAmelCase__    :							List[Any]									=		["prompt"]
     UpperCAmelCase__    :							List[str]									=		[
         "num_images_per_prompt",
         "num_inference_steps",
         "generator",
         "latents",
         "guidance_scale",
         "frame_size",
         "output_type",
         "return_dict",
     ]
     UpperCAmelCase__    :							Dict									=		False





     @property
     def 		UpperCAmelCase(self				:   Optional[Any]       )						->						str:
         return 3_2





     @property
     def 		UpperCAmelCase(self				:   Dict       )						->						Optional[Any]:
         return 3_2





     @property
     def 		UpperCAmelCase(self				:   Union[str, Any]       )						->						Union[str, Any]:
         return self.time_input_dim * 4





     @property
     def 		UpperCAmelCase(self				:   Dict       )						->						Union[str, Any]:
         return 8





     @property
     def 		UpperCAmelCase(self				:   List[str]       )						->						Dict:
         snake_case								= CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"       )
         return tokenizer





     @property
     def 		UpperCAmelCase(self				:   Dict       )						->						str:
         torch.manual_seed(0       )
         snake_case								= 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=3_7    ,   layer_norm_eps=1E-05    ,   num_attention_heads=4    ,   num_hidden_layers=5    ,   pad_token_id=1    ,   vocab_size=1_0_0_0    ,   )
         return CLIPTextModelWithProjection(snake_case_       )





     @property
     def 		UpperCAmelCase(self				:   List[Any]       )						->						Union[str, Any]:
         torch.manual_seed(0       )

         snake_case								= {
             '''num_attention_heads''': 2,
             '''attention_head_dim''': 1_6,
             '''embedding_dim''': self.time_input_dim,
             '''num_embeddings''': 3_2,
             '''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								= PriorTransformer(**snake_case_       )
         return model





     @property
     def 		UpperCAmelCase(self				:   Tuple       )						->						Optional[Any]:
         torch.manual_seed(0       )

         snake_case								= {
             '''param_shapes''': (
                 (self.renderer_dim, 9_3),
                 (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''': 1_2,
             '''background''': (
                 0.1,
                 0.1,
                 0.1,
             ),
         }
         snake_case								= ShapERenderer(**snake_case_       )
         return model





     def 		UpperCAmelCase(self				:   Any       )						->						List[Any]:
         snake_case								= self.dummy_prior
         snake_case								= self.dummy_text_encoder
         snake_case								= self.dummy_tokenizer
         snake_case								= self.dummy_renderer

         snake_case								= HeunDiscreteScheduler(
             beta_schedule="exp"    ,   num_train_timesteps=1_0_2_4    ,   prediction_type="sample"    ,   use_karras_sigmas=snake_case_    ,   clip_sample=snake_case_    ,   clip_sample_range=1.0    ,   )
         snake_case								= {
             '''prior''': prior,
             '''text_encoder''': text_encoder,
             '''tokenizer''': tokenizer,
             '''renderer''': renderer,
             '''scheduler''': scheduler,
         }

         return components





     def 		UpperCAmelCase(self				:   str    ,   _A				:   Dict    ,   _A				:   Optional[int]=0       )						->						Optional[int]:
         if str(snake_case_       ).startswith("mps"       ):
             snake_case								= torch.manual_seed(snake_case_       )
         else:
             snake_case								= torch.Generator(device=snake_case_       ).manual_seed(snake_case_       )
         snake_case								= {
             '''prompt''': '''horse''',
             '''generator''': generator,
             '''num_inference_steps''': 1,
             '''frame_size''': 3_2,
             '''output_type''': '''np''',
         }
         return inputs





     def 		UpperCAmelCase(self				:   int       )						->						Dict:
         snake_case								= '''cpu'''

         snake_case								= self.get_dummy_components()

         snake_case								= self.pipeline_class(**snake_case_       )
         snake_case								= pipe.to(snake_case_       )

         pipe.set_progress_bar_config(disable=snake_case_       )

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

         assert image.shape == (2_0, 3_2, 3_2, 3)

         snake_case								= np.array(
             [
                 0.00_03_92_16,
                 0.00_03_92_16,
                 0.00_03_92_16,
                 0.00_03_92_16,
                 0.00_03_92_16,
                 0.00_03_92_16,
                 0.00_03_92_16,
                 0.00_03_92_16,
                 0.00_03_92_16,
             ]       )

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





     def 		UpperCAmelCase(self				:   Any       )						->						Tuple:
         # 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				:   Tuple       )						->						str:
         snake_case								= torch_device == '''cpu'''
         snake_case								= True

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





     def 		UpperCAmelCase(self				:   Tuple       )						->						Optional[Any]:
         snake_case								= self.get_dummy_components()
         snake_case								= self.pipeline_class(**snake_case_       )
         snake_case								= pipe.to(snake_case_       )
         pipe.set_progress_bar_config(disable=snake_case_       )

         snake_case								= 1
         snake_case								= 2

         snake_case								= self.get_dummy_inputs(snake_case_       )

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

         snake_case								= pipe(**snake_case_    ,   num_images_per_prompt=snake_case_       )[0]

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







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





     def 		UpperCAmelCase(self				:   Dict       )						->						List[Any]:
         # clean up the VRAM after each test
         super().tearDown()
         gc.collect()
         torch.cuda.empty_cache()





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

         snake_case								= torch.Generator(device=snake_case_       ).manual_seed(0       )

         snake_case								= pipe(
             "a shark"    ,   generator=snake_case_    ,   guidance_scale=15.0    ,   num_inference_steps=6_4    ,   frame_size=6_4    ,   output_type="np"    ,   ).images[0]

         assert images.shape == (2_0, 6_4, 6_4, 3)

         assert_mean_pixel_difference(snake_case_    ,   snake_case_       )



import warnings
from typing import List, Optional, Union

from ...processing_utils import ProcessorMixin
from ...tokenization_utils_base import BatchEncoding, PaddingStrategy, PreTokenizedInput, TextInput, TruncationStrategy
from ...utils import TensorType







class 					lowerCamelCase		(  A_  ):
     UpperCAmelCase__    :							Dict									=		["image_processor", "tokenizer"]
     UpperCAmelCase__    :							Dict									=		"LayoutLMv2ImageProcessor"
     UpperCAmelCase__    :							Optional[Any]									=		("LayoutXLMTokenizer", "LayoutXLMTokenizerFast")





     def __init__(self				:   str    ,   _A				:   Any=None    ,   _A				:   Tuple=None    ,   **_A				:   Optional[Any]       )						->						Optional[int]:
         if "feature_extractor" in kwargs:
             warnings.warn(
                 "The `feature_extractor` argument is deprecated and will be removed in v5, use `image_processor`"
                 " instead."    ,   _A    ,   )
             snake_case								= kwargs.pop("feature_extractor"       )

         snake_case								= image_processor if image_processor is not None else feature_extractor
         if image_processor is None:
             raise ValueError("You need to specify an `image_processor`."       )
         if tokenizer is None:
             raise ValueError("You need to specify a `tokenizer`."       )

         super().__init__(_A    ,   _A       )





     def __call__(self				:   int    ,   _A				:   List[str]    ,   _A				:   Union[TextInput, PreTokenizedInput, List[TextInput], List[PreTokenizedInput]] = None    ,   _A				:   Optional[Union[PreTokenizedInput, List[PreTokenizedInput]]] = None    ,   _A				:   Union[List[List[int]], List[List[List[int]]]] = None    ,   _A				:   Optional[Union[List[int], List[List[int]]]] = None    ,   _A				:   bool = True    ,   _A				:   Union[bool, str, PaddingStrategy] = False    ,   _A				:   Union[bool, str, TruncationStrategy] = None    ,   _A				:   Optional[int] = None    ,   _A				:   int = 0    ,   _A				:   Optional[int] = None    ,   _A				:   Optional[bool] = None    ,   _A				:   Optional[bool] = None    ,   _A				:   bool = False    ,   _A				:   bool = False    ,   _A				:   bool = False    ,   _A				:   bool = False    ,   _A				:   bool = True    ,   _A				:   Optional[Union[str, TensorType]] = None    ,   **_A				:   Dict    ,   )						->						BatchEncoding:
         # verify input
         if self.image_processor.apply_ocr and (boxes is not None):
             raise ValueError(
                 "You cannot provide bounding boxes "
                 "if you initialized the image processor with apply_ocr set to True."       )

         if self.image_processor.apply_ocr and (word_labels is not None):
             raise ValueError(
                 "You cannot provide word labels if you initialized the image processor with apply_ocr set to True."       )

         if return_overflowing_tokens is True and return_offsets_mapping is False:
             raise ValueError("You cannot return overflowing tokens without returning the offsets mapping."       )

         # first, apply the image processor
         snake_case								= self.image_processor(images=_A    ,   return_tensors=_A       )

         # second, apply the tokenizer
         if text is not None and self.image_processor.apply_ocr and text_pair is None:
             if isinstance(_A    ,   _A       ):
                 snake_case								= [text]  # add batch dimension (as the image processor always adds a batch dimension)
             snake_case								= features["words"]

         snake_case								= self.tokenizer(
             text=text if text is not None else features["words"]    ,   text_pair=text_pair if text_pair is not None else None    ,   boxes=boxes if boxes is not None else features["boxes"]    ,   word_labels=_A    ,   add_special_tokens=_A    ,   padding=_A    ,   truncation=_A    ,   max_length=_A    ,   stride=_A    ,   pad_to_multiple_of=_A    ,   return_token_type_ids=_A    ,   return_attention_mask=_A    ,   return_overflowing_tokens=_A    ,   return_special_tokens_mask=_A    ,   return_offsets_mapping=_A    ,   return_length=_A    ,   verbose=_A    ,   return_tensors=_A    ,   **_A    ,   )

         # add pixel values
         snake_case								= features.pop("pixel_values"       )
         if return_overflowing_tokens is True:
             snake_case								= self.get_overflowing_images(_A    ,   encoded_inputs["overflow_to_sample_mapping"]       )
         snake_case								= images

         return encoded_inputs





     def 		UpperCAmelCase(self				:   Dict    ,   _A				:   Dict    ,   _A				:   List[str]       )						->						Optional[int]:
         # in case there's an overflow, ensure each `input_ids` sample is mapped to its corresponding image
         snake_case								= []
         for sample_idx in overflow_to_sample_mapping:
             images_with_overflow.append(images[sample_idx]       )

         if len(_A       ) != len(_A       ):
             raise ValueError(
                 "Expected length of images to be the same as the length of `overflow_to_sample_mapping`, but got"
                 f' {len(_A       )} and {len(_A       )}'       )

         return images_with_overflow





     def 		UpperCAmelCase(self				:   Tuple    ,   *_A				:   int    ,   **_A				:   Dict       )						->						str:
         return self.tokenizer.batch_decode(*_A    ,   **_A       )





     def 		UpperCAmelCase(self				:   str    ,   *_A				:   List[Any]    ,   **_A				:   List[Any]       )						->						Optional[Any]:
         return self.tokenizer.decode(*_A    ,   **_A       )





     @property
     def 		UpperCAmelCase(self				:   Tuple       )						->						Optional[int]:
         return ["input_ids", "bbox", "attention_mask", "image"]





     @property
     def 		UpperCAmelCase(self				:   List[Any]       )						->						int:
         warnings.warn(
             "`feature_extractor_class` is deprecated and will be removed in v5. Use `image_processor_class` instead."    ,   _A    ,   )
         return self.image_processor_class





     @property
     def 		UpperCAmelCase(self				:   Dict       )						->						Union[str, Any]:
         warnings.warn(
             "`feature_extractor` is deprecated and will be removed in v5. Use `image_processor` instead."    ,   _A    ,   )
         return self.image_processor

