



'''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  SCREAMING_SNAKE_CASE (  a__   ,						unittest.TestCase		):






    """simple docstring"""





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






    @property
    def  A       (  self      :					int ):


       """simple docstring"""







       return 3_2






    @property
    def  A       (  self      :					Tuple ):


       """simple docstring"""







       return 3_2






    @property
    def  A       (  self      :					Union[str, Any] ):


       """simple docstring"""







       return self.time_input_dim * 4






    @property
    def  A       (  self      :					List[str] ):


       """simple docstring"""







       return 8






    @property
    def  A       (  self      :					List[Any] ):


       """simple docstring"""







       UpperCamelCase					=  CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip' )
       return tokenizer






    @property
    def  A       (  self      :					Union[str, Any] ):


       """simple docstring"""







       torch.manual_seed(0 )
       UpperCamelCase					=  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-0_5  ,						num_attention_heads=4  ,						num_hidden_layers=5  ,						pad_token_id=1  ,						vocab_size=1_0_0_0  ,						)
       return CLIPTextModelWithProjection(SCREAMING_SNAKE_CASE_ )






    @property
    def  A       (  self      :					Optional[int] ):


       """simple docstring"""







       torch.manual_seed(0 )

       UpperCamelCase					=  {
           '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,
       }

       UpperCamelCase					=  PriorTransformer(**SCREAMING_SNAKE_CASE_ )
       return model






    @property
    def  A       (  self      :					Dict ):


       """simple docstring"""







       torch.manual_seed(0 )

       UpperCamelCase					=  {
           '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,
           ),
       }
       UpperCamelCase					=  ShapERenderer(**SCREAMING_SNAKE_CASE_ )
       return model






    def  A       (  self      :					Tuple ):


       """simple docstring"""







       UpperCamelCase					=  self.dummy_prior
       UpperCamelCase					=  self.dummy_text_encoder
       UpperCamelCase					=  self.dummy_tokenizer
       UpperCamelCase					=  self.dummy_renderer

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

       return components






    def  A       (  self      :					Tuple  ,						UpperCamelCase__      :					List[Any]  ,						UpperCamelCase__      :					Optional[Any]=0 ):


       """simple docstring"""







       if str(SCREAMING_SNAKE_CASE_ ).startswith('mps' ):
          UpperCamelCase					=  torch.manual_seed(SCREAMING_SNAKE_CASE_ )
       else:
          UpperCamelCase					=  torch.Generator(device=SCREAMING_SNAKE_CASE_ ).manual_seed(SCREAMING_SNAKE_CASE_ )
       UpperCamelCase					=  {
           'prompt': 'horse',
           'generator': generator,
           'num_inference_steps': 1,
           'frame_size': 3_2,
           'output_type': 'np',
       }
       return inputs






    def  A       (  self      :					List[Any] ):


       """simple docstring"""







       UpperCamelCase					=  'cpu'

       UpperCamelCase					=  self.get_dummy_components()

       UpperCamelCase					=  self.pipeline_class(**SCREAMING_SNAKE_CASE_ )
       UpperCamelCase					=  pipe.to(SCREAMING_SNAKE_CASE_ )

       pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE_ )

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

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

       UpperCamelCase					=  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  A       (  self      :					Tuple ):


       """simple docstring"""







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






    def  A       (  self      :					Any ):


       """simple docstring"""







       UpperCamelCase					=  torch_device == 'cpu'
       UpperCamelCase					=  True

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






    def  A       (  self      :					Optional[Any] ):


       """simple docstring"""







       UpperCamelCase					=  self.get_dummy_components()
       UpperCamelCase					=  self.pipeline_class(**SCREAMING_SNAKE_CASE_ )
       UpperCamelCase					=  pipe.to(SCREAMING_SNAKE_CASE_ )
       pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE_ )

       UpperCamelCase					=  1
       UpperCamelCase					=  2

       UpperCamelCase					=  self.get_dummy_inputs(SCREAMING_SNAKE_CASE_ )

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

       UpperCamelCase					=  pipe(**SCREAMING_SNAKE_CASE_  ,						num_images_per_prompt=SCREAMING_SNAKE_CASE_ )[0]

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

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






    """simple docstring"""





    def  A       (  self      :					int ):


       """simple docstring"""







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






    def  A       (  self      :					str ):


       """simple docstring"""







       UpperCamelCase					=  load_numpy(
           'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
           '/shap_e/test_shap_e_np_out.npy' )
       UpperCamelCase					=  ShapEPipeline.from_pretrained('openai/shap-e' )
       UpperCamelCase					=  pipe.to(SCREAMING_SNAKE_CASE_ )
       pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE_ )

       UpperCamelCase					=  torch.Generator(device=SCREAMING_SNAKE_CASE_ ).manual_seed(0 )

       UpperCamelCase					=  pipe(
           'a shark'  ,						generator=SCREAMING_SNAKE_CASE_  ,						guidance_scale=1_5.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(SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_ )






from typing import Optional

import pyspark

from .. import Features, NamedSplit
from ..download import DownloadMode
from ..packaged_modules.spark.spark import Spark
from .abc import AbstractDatasetReader



class 				lowerCAmelCase_		(  a__							):
       def __init__(			self,			SCREAMING_SNAKE_CASE_,			SCREAMING_SNAKE_CASE_ = None,			SCREAMING_SNAKE_CASE_ = None,			SCREAMING_SNAKE_CASE_ = True,			SCREAMING_SNAKE_CASE_ = None,			SCREAMING_SNAKE_CASE_ = False,			SCREAMING_SNAKE_CASE_ = None,			SCREAMING_SNAKE_CASE_ = True,			SCREAMING_SNAKE_CASE_ = "arrow",			**SCREAMING_SNAKE_CASE_,			)     ->		Optional[int]:
          super().__init__(
              split=SCREAMING_SNAKE_CASE_,			features=SCREAMING_SNAKE_CASE_,			cache_dir=SCREAMING_SNAKE_CASE_,			keep_in_memory=SCREAMING_SNAKE_CASE_,			streaming=SCREAMING_SNAKE_CASE_,			**SCREAMING_SNAKE_CASE_,			)
          UpperCamelCase    :  List[str]							=	load_from_cache_file
          UpperCamelCase    :  List[str]							=	file_format
          UpperCamelCase    :  Optional[int]							=	Spark(
              df=SCREAMING_SNAKE_CASE_,			features=SCREAMING_SNAKE_CASE_,			cache_dir=SCREAMING_SNAKE_CASE_,			working_dir=SCREAMING_SNAKE_CASE_,			**SCREAMING_SNAKE_CASE_,			)






       def 					snake_case_       (			self       )     ->		int:
          if self.streaming:
             return self.builder.as_streaming_dataset(split=self.split       )
          UpperCamelCase    :  Union[str, Any]							=	None if self._load_from_cache_file else DownloadMode.FORCE_REDOWNLOAD
          self.builder.download_and_prepare(
              download_mode=SCREAMING_SNAKE_CASE_,			file_format=self._file_format,			)
          return self.builder.as_dataset(split=self.split       )



