

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		(  __magic_name__   ,     unittest.TestCase):
  SCREAMING_SNAKE_CASE    :      str       =  ShapEPipeline
  SCREAMING_SNAKE_CASE    :      str       =  ['''prompt''']
  SCREAMING_SNAKE_CASE    :      Optional[int]       =  ['''prompt''']
  SCREAMING_SNAKE_CASE    :      Dict       =  [
      '''num_images_per_prompt''',
      '''num_inference_steps''',
      '''generator''',
      '''latents''',
      '''guidance_scale''',
      '''frame_size''',
      '''output_type''',
      '''return_dict''',
  ]
  SCREAMING_SNAKE_CASE    :      Union[str, Any]       =  False

  @property
  def  UpperCAmelCase    (				self					):





        """simple docstring"""







        return 32

  @property
  def  UpperCAmelCase    (				self					):





        """simple docstring"""







        return 32

  @property
  def  UpperCAmelCase    (				self					):





        """simple docstring"""







        return self.time_input_dim * 4

  @property
  def  UpperCAmelCase    (				self					):





        """simple docstring"""







        return 8

  @property
  def  UpperCAmelCase    (				self					):





        """simple docstring"""







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

  @property
  def  UpperCAmelCase    (				self					):





        """simple docstring"""







        torch.manual_seed(0					)
        SCREAMING_SNAKE_CASE_				:  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(_SCREAMING_SNAKE_CASE					)

  @property
  def  UpperCAmelCase    (				self					):





        """simple docstring"""







        torch.manual_seed(0					)

        SCREAMING_SNAKE_CASE_				:  Dict           =     {
            '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,
        }

        SCREAMING_SNAKE_CASE_				:  Tuple           =     PriorTransformer(**_SCREAMING_SNAKE_CASE					)
        return model

  @property
  def  UpperCAmelCase    (				self					):





        """simple docstring"""







        torch.manual_seed(0					)

        SCREAMING_SNAKE_CASE_				:  Optional[int]           =     {
            '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,
            ),
        }
        SCREAMING_SNAKE_CASE_				:  List[str]           =     ShapERenderer(**_SCREAMING_SNAKE_CASE					)
        return model

  def  UpperCAmelCase    (				self					):





        """simple docstring"""







        SCREAMING_SNAKE_CASE_				:  List[Any]           =     self.dummy_prior
        SCREAMING_SNAKE_CASE_				:  List[str]           =     self.dummy_text_encoder
        SCREAMING_SNAKE_CASE_				:  Tuple           =     self.dummy_tokenizer
        SCREAMING_SNAKE_CASE_				:  Optional[Any]           =     self.dummy_renderer

        SCREAMING_SNAKE_CASE_				:  str           =     HeunDiscreteScheduler(
            beta_schedule='exp' ,					num_train_timesteps=1024 ,					prediction_type='sample' ,					use_karras_sigmas=_SCREAMING_SNAKE_CASE ,					clip_sample=_SCREAMING_SNAKE_CASE ,					clip_sample_range=1.0 ,					)
        SCREAMING_SNAKE_CASE_				:  Union[str, Any]           =     {
            'prior': prior,
            'text_encoder': text_encoder,
            'tokenizer': tokenizer,
            'renderer': renderer,
            'scheduler': scheduler,
        }

        return components

  def  UpperCAmelCase    (				self ,					_SCREAMING_SNAKE_CASE ,					_SCREAMING_SNAKE_CASE=0					):





        """simple docstring"""







        if str(_SCREAMING_SNAKE_CASE					).startswith('mps'					):
              SCREAMING_SNAKE_CASE_				:  Any           =     torch.manual_seed(_SCREAMING_SNAKE_CASE					)
        else:
              SCREAMING_SNAKE_CASE_				:  str           =     torch.Generator(device=_SCREAMING_SNAKE_CASE					).manual_seed(_SCREAMING_SNAKE_CASE					)
        SCREAMING_SNAKE_CASE_				:  List[Any]           =     {
            'prompt': 'horse',
            'generator': generator,
            'num_inference_steps': 1,
            'frame_size': 32,
            'output_type': 'np',
        }
        return inputs

  def  UpperCAmelCase    (				self					):





        """simple docstring"""







        SCREAMING_SNAKE_CASE_				:  Dict           =     'cpu'

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

        SCREAMING_SNAKE_CASE_				:  Any           =     self.pipeline_class(**_SCREAMING_SNAKE_CASE					)
        SCREAMING_SNAKE_CASE_				:  Tuple           =     pipe.to(_SCREAMING_SNAKE_CASE					)

        pipe.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE					)

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

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

        SCREAMING_SNAKE_CASE_				:  Any           =     np.array(
            [
                0.00039216,
                0.00039216,
                0.00039216,
                0.00039216,
                0.00039216,
                0.00039216,
                0.00039216,
                0.00039216,
                0.00039216,
            ]					)

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

  def  UpperCAmelCase    (				self					):





        """simple docstring"""







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

  def  UpperCAmelCase    (				self					):





        """simple docstring"""







        SCREAMING_SNAKE_CASE_				:  str           =     torch_device == 'cpu'
        SCREAMING_SNAKE_CASE_				:  Tuple           =     True

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



  def  UpperCAmelCase    (				self					):





        """simple docstring"""







        SCREAMING_SNAKE_CASE_				:  Any           =     self.get_dummy_components()
        SCREAMING_SNAKE_CASE_				:  List[str]           =     self.pipeline_class(**_SCREAMING_SNAKE_CASE					)
        SCREAMING_SNAKE_CASE_				:  List[str]           =     pipe.to(_SCREAMING_SNAKE_CASE					)
        pipe.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE					)

        SCREAMING_SNAKE_CASE_				:  Tuple           =     1
        SCREAMING_SNAKE_CASE_				:  List[Any]           =     2

        SCREAMING_SNAKE_CASE_				:  Any           =     self.get_dummy_inputs(_SCREAMING_SNAKE_CASE					)

        for key in inputs.keys():
              if key in self.batch_params:
                    SCREAMING_SNAKE_CASE_				:  Union[str, Any]           =     batch_size * [inputs[key]]

        SCREAMING_SNAKE_CASE_				:  str           =     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  _A		(  unittest.TestCase):

  def  UpperCAmelCase    (				self					):





        """simple docstring"""







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



  def  UpperCAmelCase    (				self					):





        """simple docstring"""







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

        SCREAMING_SNAKE_CASE_				:  Any           =     torch.Generator(device=_SCREAMING_SNAKE_CASE					).manual_seed(0					)

        SCREAMING_SNAKE_CASE_				:  str           =     pipe(
            'a shark' ,					generator=_SCREAMING_SNAKE_CASE ,					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(_SCREAMING_SNAKE_CASE ,					_SCREAMING_SNAKE_CASE					)









from __future__ import annotations

import math

lowerCAmelCase     :   List[str]						     =  '2020.9.26'
lowerCAmelCase     :   List[Any]						     =  'xcodz-dot, cclaus, dhruvmanila'

def    A_  (							a					,   a					,   a					,   a					,   a  ):

      """simple docstring"""



      if not all(isinstance(a					,   (float, int)  ) for val in locals().values()  ):
            SCREAMING_SNAKE_CASE_				:  Optional[Any]           =     f"Input values must either be float or int: {list(locals().values()  )}"
            raise TypeError(a  )
      SCREAMING_SNAKE_CASE_				:  Optional[Any]           =     ((x * distance) / (z + distance)) * scale
      SCREAMING_SNAKE_CASE_				:  int           =     ((y * distance) / (z + distance)) * scale
      return projected_x, projected_y

def    A_  (							a					,   a					,   a					,   a					,   a  ):

      """simple docstring"""



      if not isinstance(a					,   a  ):
            raise TypeError('Axis must be a str'  )
      SCREAMING_SNAKE_CASE_				:  Any           =     locals()
      del input_variables["axis"]
      if not all(isinstance(a					,   (float, int)  ) for val in input_variables.values()  ):
            SCREAMING_SNAKE_CASE_				:  Optional[Any]           =     (
                'Input values except axis must either be float or int: '
                f"{list(input_variables.values()  )}"
            )
            raise TypeError(a  )
      SCREAMING_SNAKE_CASE_				:  List[Any]           =     (angle % 3_6_0) / 4_5_0 * 1_8_0 / math.pi
      if axis == "z":
            SCREAMING_SNAKE_CASE_				:  List[str]           =     x * math.cos(a  ) - y * math.sin(a  )
            SCREAMING_SNAKE_CASE_				:  List[Any]           =     y * math.cos(a  ) + x * math.sin(a  )
            SCREAMING_SNAKE_CASE_				:  Dict           =     z
      elif axis == "x":
            SCREAMING_SNAKE_CASE_				:  Any           =     y * math.cos(a  ) - z * math.sin(a  )
            SCREAMING_SNAKE_CASE_				:  Optional[Any]           =     z * math.cos(a  ) + y * math.sin(a  )
            SCREAMING_SNAKE_CASE_				:  str           =     x
      elif axis == "y":
            SCREAMING_SNAKE_CASE_				:  str           =     x * math.cos(a  ) - z * math.sin(a  )
            SCREAMING_SNAKE_CASE_				:  List[str]           =     z * math.cos(a  ) + x * math.sin(a  )
            SCREAMING_SNAKE_CASE_				:  Any           =     y
      else:
            raise ValueError('not a valid axis, choose one of \'x\', \'y\', \'z\''  )

      return new_x, new_y, new_z


if __name__ == "__main__":
    import doctest

    doctest.testmod()
    print(F'{convert_to_ad(1.0, 2.0, 3.0, 10.0, 10.0) = }')
    print(F'{rotate(1.0, 2.0, 3.0, "y", 90.0) = }')






