






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

       snake_case__  :    Dict							       ={
           """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__  :    str							       =PriorTransformer(**a   )
       return model
      @property
      def 						lowercase__					(     self   ):
       torch.manual_seed(0   )

       snake_case__  :    str							       ={
           """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__  :    Optional[int]							       =ShapERenderer(**a   )
       return model
      def 						lowercase__					(     self   ):
       snake_case__  :    Dict							       =self.dummy_prior
       snake_case__  :    Any							       =self.dummy_text_encoder
       snake_case__  :    Optional[Any]							       =self.dummy_tokenizer
       snake_case__  :    Optional[int]							       =self.dummy_renderer

       snake_case__  :    List[Any]							       =HeunDiscreteScheduler(
           beta_schedule="""exp"""				,			num_train_timesteps=1_0_2_4				,			prediction_type="""sample"""				,			use_karras_sigmas=a				,			clip_sample=a				,			clip_sample_range=1.0				,			)
       snake_case__  :    Tuple							       ={
           """prior""": prior,
           """text_encoder""": text_encoder,
           """tokenizer""": tokenizer,
           """renderer""": renderer,
           """scheduler""": scheduler,
       }

       return components
      def 						lowercase__					(     self				,			a				,			a=0   ):
       if str(a   ).startswith("""mps"""   ):
        snake_case__  :    List[Any]							       =torch.manual_seed(a   )
       else:
        snake_case__  :    Any							       =torch.Generator(device=a   ).manual_seed(a   )
       snake_case__  :    str							       ={
           """prompt""": """horse""",
           """generator""": generator,
           """num_inference_steps""": 1,
           """frame_size""": 3_2,
           """output_type""": """np""",
       }
       return inputs
      def 						lowercase__					(     self   ):
       snake_case__  :    Any							       ="""cpu"""

       snake_case__  :    Union[str, Any]							       =self.get_dummy_components()

       snake_case__  :    Union[str, Any]							       =self.pipeline_class(**a   )
       snake_case__  :    int							       =pipe.to(a   )

       pipe.set_progress_bar_config(disable=a   )

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

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

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

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



      def 						lowercase__					(     self   ):
       snake_case__  :    int							       =self.get_dummy_components()
       snake_case__  :    Optional[Any]							       =self.pipeline_class(**a   )
       snake_case__  :    List[Any]							       =pipe.to(a   )
       pipe.set_progress_bar_config(disable=a   )

       snake_case__  :    Dict							       =1
       snake_case__  :    Dict							       =2

       snake_case__  :    Dict							       =self.get_dummy_inputs(a   )

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

       snake_case__  :    Dict							       =pipe(**a				,			num_images_per_prompt=a   )[0]

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







@slow
@require_torch_gpu
class  _lowercase   (   unittest.TestCase						):
      def 						lowercase__					(     self   ):
       # clean up the VRAM after each test
       super().tearDown()
       gc.collect()
       torch.cuda.empty_cache()



      def 						lowercase__					(     self   ):
       snake_case__  :    str							       =load_numpy(
           """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"""
           """/shap_e/test_shap_e_np_out.npy"""   )
       snake_case__  :    Dict							       =ShapEPipeline.from_pretrained("""openai/shap-e"""   )
       snake_case__  :    Optional[Any]							       =pipe.to(a   )
       pipe.set_progress_bar_config(disable=a   )

       snake_case__  :    Optional[Any]							       =torch.Generator(device=a   ).manual_seed(0   )

       snake_case__  :    Union[str, Any]							       =pipe(
           """a shark"""				,			generator=a				,			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(a				,			a   )













__lowerCamelCase	:		dict[str, float]										=   {
    "km/h": 1.0,
    "m/s": 3.6,
    "mph": 1.60_93_44,
    "knot": 1.8_52,
}

__lowerCamelCase	:		dict[str, float]										=   {
    "km/h": 1.0,
    "m/s": 0.2_77_77_77_78,
    "mph": 0.6_21_37_11_92,
    "knot": 0.5_39_95_68_03,
}







def 			A__							(      _a  :					float		,  _a  :					str		,  _a  :					str	):





 '''simple docstring'''







 if unit_to not in speed_chart or unit_from not in speed_chart_inverse:
  snake_case__  :    Tuple							       =(
      f"Incorrect 'from_type' or 'to_type' value: {unit_from!r}, {unit_to!r}\n"
      f"Valid values are: {', '.join(_a	)}"
  )
  raise ValueError(_a	)
 return round(speed * speed_chart[unit_from] * speed_chart_inverse[unit_to]		,  3	)


if __name__ == "__main__":
       import doctest

       doctest.testmod()





