
"""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       __magic_name__							(							SCREAMING_SNAKE_CASE__ ,       unittest.TestCase			):
       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 				lowercase_				(  self							)  ->						Dict:





         """simple docstring"""




         return 32
       @property
       def 				lowercase_				(  self							)  ->						Any:





         """simple docstring"""




         return 32
       @property
       def 				lowercase_				(  self							)  ->						Tuple:





         """simple docstring"""




         return self.time_input_dim * 4
       @property
       def 				lowercase_				(  self							)  ->						str:





         """simple docstring"""




         return 8
       @property
       def 				lowercase_				(  self							)  ->						List[str]:





         """simple docstring"""




         _lowercase:      Tuple   			=  CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''							)
         return tokenizer
       @property
       def 				lowercase_				(  self							)  ->						str:





         """simple docstring"""




         torch.manual_seed(0							)
         _lowercase:      Any   			=  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(A_							)
       @property
       def 				lowercase_				(  self							)  ->						Any:





         """simple docstring"""




         torch.manual_seed(0							)

         _lowercase:      Tuple   			=  {
             '''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,
         }

         _lowercase:      int   			=  PriorTransformer(**A_							)
         return model
       @property
       def 				lowercase_				(  self							)  ->						Any:





         """simple docstring"""




         torch.manual_seed(0							)

         _lowercase:      Dict   			=  {
             '''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,
             ),
         }
         _lowercase:      List[Any]   			=  ShapERenderer(**A_							)
         return model
       def 				lowercase_				(  self							)  ->						Dict:





         """simple docstring"""




         _lowercase:      Optional[Any]   			=  self.dummy_prior
         _lowercase:      Tuple   			=  self.dummy_text_encoder
         _lowercase:      Union[str, Any]   			=  self.dummy_tokenizer
         _lowercase:      Optional[int]   			=  self.dummy_renderer

         _lowercase:      Dict   			=  HeunDiscreteScheduler(
             beta_schedule='''exp'''			,					num_train_timesteps=1024			,					prediction_type='''sample'''			,					use_karras_sigmas=A_			,					clip_sample=A_			,					clip_sample_range=1.0			,					)
         _lowercase:      List[str]   			=  {
             '''prior''': prior,
             '''text_encoder''': text_encoder,
             '''tokenizer''': tokenizer,
             '''renderer''': renderer,
             '''scheduler''': scheduler,
         }

         return components
       def 				lowercase_				(  self			,					A_			,					A_=0							)  ->						int:





         """simple docstring"""




         if str(A_							).startswith('''mps'''							):
           _lowercase:      int   			=  torch.manual_seed(A_							)
         else:
           _lowercase:      int   			=  torch.Generator(device=A_							).manual_seed(A_							)
         _lowercase:      Any   			=  {
             '''prompt''': '''horse''',
             '''generator''': generator,
             '''num_inference_steps''': 1,
             '''frame_size''': 32,
             '''output_type''': '''np''',
         }
         return inputs
       def 				lowercase_				(  self							)  ->						int:





         """simple docstring"""




         _lowercase:      int   			=  '''cpu'''

         _lowercase:      Optional[Any]   			=  self.get_dummy_components()

         _lowercase:      str   			=  self.pipeline_class(**A_							)
         _lowercase:      Any   			=  pipe.to(A_							)

         pipe.set_progress_bar_config(disable=A_							)

         _lowercase:      Union[str, Any]   			=  pipe(**self.get_dummy_inputs(A_							)							)
         _lowercase:      Any   			=  output.images[0]
         _lowercase:      List[Any]   			=  image[0, -3:, -3:, -1]

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

         _lowercase:      List[Any]   			=  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 				lowercase_				(  self							)  ->						Union[str, Any]:





         """simple docstring"""




         self._test_inference_batch_consistent(batch_sizes=[1, 2]							)
       def 				lowercase_				(  self							)  ->						List[Any]:





         """simple docstring"""




         _lowercase:      Dict   			=  torch_device == '''cpu'''
         _lowercase:      Tuple   			=  True

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



       def 				lowercase_				(  self							)  ->						Optional[Any]:





         """simple docstring"""




         _lowercase:      List[Any]   			=  self.get_dummy_components()
         _lowercase:      Any   			=  self.pipeline_class(**A_							)
         _lowercase:      Optional[Any]   			=  pipe.to(A_							)
         pipe.set_progress_bar_config(disable=A_							)

         _lowercase:      str   			=  1
         _lowercase:      Optional[Any]   			=  2

         _lowercase:      str   			=  self.get_dummy_inputs(A_							)

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

         _lowercase:      List[Any]   			=  pipe(**A_			,					num_images_per_prompt=A_							)[0]

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


@slow
@require_torch_gpu
class       __magic_name__							(							unittest.TestCase			):
       def 				lowercase_				(  self							)  ->						Optional[Any]:





         """simple docstring"""




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



       def 				lowercase_				(  self							)  ->						int:





         """simple docstring"""




         _lowercase:      Tuple   			=  load_numpy(
             '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
             '''/shap_e/test_shap_e_np_out.npy'''							)
         _lowercase:      Optional[int]   			=  ShapEPipeline.from_pretrained('''openai/shap-e'''							)
         _lowercase:      int   			=  pipe.to(A_							)
         pipe.set_progress_bar_config(disable=A_							)

         _lowercase:      Dict   			=  torch.Generator(device=A_							).manual_seed(0							)

         _lowercase:      str   			=  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"""




import numpy as np
import torch
from torch.utils.data import Dataset

from utils import logger


class       __magic_name__							(							SCREAMING_SNAKE_CASE__			):
       def __init__(  self			,					A_			,					A_							)  ->						List[str]:





         """simple docstring"""




         _lowercase:      List[str]   			=  params

         _lowercase:      str   			=  np.array(A_							)
         _lowercase:      Optional[int]   			=  np.array([len(A_							) for t in data]							)

         self.check()
         self.remove_long_sequences()
         self.remove_empty_sequences()
         self.remove_unknown_sequences()
         self.check()
         self.print_statistics()
       def __getitem__(  self			,					A_							)  ->						Optional[Any]:





         """simple docstring"""




         return (self.token_ids[index], self.lengths[index])
       def __len__(  self							)  ->						List[str]:





         """simple docstring"""




         return len(self.lengths							)
       def 				lowercase_				(  self							)  ->						Dict:





         """simple docstring"""




         assert len(self.token_ids							) == len(self.lengths							)
         assert all(self.lengths[i] == len(self.token_ids[i]							) for i in range(len(self.lengths							)							)							)
       def 				lowercase_				(  self							)  ->						str:





         """simple docstring"""




         _lowercase:      Tuple   			=  self.params.max_model_input_size
         _lowercase:      Tuple   			=  self.lengths > max_len
         logger.info(f'''Splitting {sum(A_							)} too long sequences.'''							)

         def divide_chunks(A_			,					A_							):
           return [l[i : i + n] for i in range(0			,					len(A_							)			,					A_							)]

         _lowercase:      Dict   			=  []
         _lowercase:      Union[str, Any]   			=  []
         if self.params.mlm:
           _lowercase			,					_lowercase:      int   			=  self.params.special_tok_ids['''cls_token'''], self.params.special_tok_ids['''sep_token''']
         else:
           _lowercase			,					_lowercase:      Dict   			=  self.params.special_tok_ids['''bos_token'''], self.params.special_tok_ids['''eos_token''']

         for seq_, len_ in zip(self.token_ids			,					self.lengths							):
           assert (seq_[0] == cls_id) and (seq_[-1] == sep_id), seq_
           if len_ <= max_len:
             new_tok_ids.append(seq_							)
             new_lengths.append(len_							)
           else:
             _lowercase:      Optional[Any]   			=  []
             for sub_s in divide_chunks(seq_			,					max_len - 2							):
               if sub_s[0] != cls_id:
                 _lowercase:      Any   			=  np.insert(A_			,					0			,					A_							)
               if sub_s[-1] != sep_id:
                 _lowercase:      Optional[int]   			=  np.insert(A_			,					len(A_							)			,					A_							)
               assert len(A_							) <= max_len
               assert (sub_s[0] == cls_id) and (sub_s[-1] == sep_id), sub_s
               sub_seqs.append(A_							)

             new_tok_ids.extend(A_							)
             new_lengths.extend([len(A_							) for l in sub_seqs]							)

         _lowercase:      Optional[Any]   			=  np.array(A_							)
         _lowercase:      List[str]   			=  np.array(A_							)
       def 				lowercase_				(  self							)  ->						Tuple:





         """simple docstring"""




         _lowercase:      List[Any]   			=  len(self							)
         _lowercase:      Optional[Any]   			=  self.lengths > 11
         _lowercase:      int   			=  self.token_ids[indices]
         _lowercase:      List[str]   			=  self.lengths[indices]
         _lowercase:      Optional[int]   			=  len(self							)
         logger.info(f'''Remove {init_size - new_size} too short (<=11 tokens) sequences.'''							)
       def 				lowercase_				(  self							)  ->						int:





         """simple docstring"""




         if "unk_token" not in self.params.special_tok_ids:
           return
         else:
           _lowercase:      Dict   			=  self.params.special_tok_ids['''unk_token''']
         _lowercase:      int   			=  len(self							)
         _lowercase:      Tuple   			=  np.array([np.count_nonzero(a == unk_token_id							) for a in self.token_ids]							)
         _lowercase:      Dict   			=  (unk_occs / self.lengths) < 0.5
         _lowercase:      str   			=  self.token_ids[indices]
         _lowercase:      Dict   			=  self.lengths[indices]
         _lowercase:      int   			=  len(self							)
         logger.info(f'''Remove {init_size - new_size} sequences with a high level of unknown tokens (50%).'''							)
       def 				lowercase_				(  self							)  ->						Optional[Any]:





         """simple docstring"""




         if not self.params.is_master:
           return
         logger.info(f'''{len(self							)} sequences'''							)
         # data_len = sum(self.lengths)
         # nb_unique_tokens = len(Counter(list(chain(*self.token_ids))))
         # logger.info(f'{data_len} tokens ({nb_unique_tokens} unique)')

         # unk_idx = self.params.special_tok_ids['unk_token']
         # nb_unknown = sum([(t==unk_idx).sum() for t in self.token_ids])
         # logger.info(f'{nb_unknown} unknown tokens (covering {100*nb_unknown/data_len:.2f}% of the data)')



       def 				lowercase_				(  self			,					A_							)  ->						Optional[int]:





         """simple docstring"""




         _lowercase:      List[str]   			=  [t[0] for t in batch]
         _lowercase:      Dict   			=  [t[1] for t in batch]
         assert len(A_							) == len(A_							)

         # Max for paddings
         _lowercase:      Tuple   			=  max(A_							)

         # Pad token ids
         if self.params.mlm:
           _lowercase:      str   			=  self.params.special_tok_ids['''pad_token''']
         else:
           _lowercase:      Union[str, Any]   			=  self.params.special_tok_ids['''unk_token''']
         _lowercase:      int   			=  [list(t.astype(A_							)							) + [pad_idx] * (max_seq_len_ - len(A_							)) for t in token_ids]
         assert len(tk_							) == len(A_							)
         assert all(len(A_							) == max_seq_len_ for t in tk_							)

         _lowercase:      str   			=  torch.tensor(tk_							)  # (bs, max_seq_len_)
         _lowercase:      int   			=  torch.tensor(A_							)  # (bs)
         return tk_t, lg_t
