




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_							(      _SCREAMING_SNAKE_CASE       ,   unittest.TestCase					):


 '''simple docstring'''







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



 @property
 def       lowerCAmelCase_  (							self  :		Tuple						):
       return 32



 @property
 def       lowerCAmelCase_  (							self  :		Any						):
       return 32



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



 @property
 def       lowerCAmelCase_  (							self  :		Dict						):
       return 8



 @property
 def       lowerCAmelCase_  (							self  :		List[Any]						):
       SCREAMING_SNAKE_CASE_							=			CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'						)
       return tokenizer



 @property
 def       lowerCAmelCase_  (							self  :		List[str]						):
       torch.manual_seed(0						)
       SCREAMING_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=37				,				layer_norm_eps=1E-05				,				num_attention_heads=4				,				num_hidden_layers=5				,				pad_token_id=1				,				vocab_size=1_000				,				)
       return CLIPTextModelWithProjection(_lowerCAmelCase						)



 @property
 def       lowerCAmelCase_  (							self  :		Dict						):
       torch.manual_seed(0						)

       SCREAMING_SNAKE_CASE_							=			{
           '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_							=			PriorTransformer(**_lowerCAmelCase						)
       return model



 @property
 def       lowerCAmelCase_  (							self  :		Optional[int]						):
       torch.manual_seed(0						)

       SCREAMING_SNAKE_CASE_							=			{
           '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_							=			ShapERenderer(**_lowerCAmelCase						)
       return model



 def       lowerCAmelCase_  (							self  :		Optional[int]						):
       SCREAMING_SNAKE_CASE_							=			self.dummy_prior
       SCREAMING_SNAKE_CASE_							=			self.dummy_text_encoder
       SCREAMING_SNAKE_CASE_							=			self.dummy_tokenizer
       SCREAMING_SNAKE_CASE_							=			self.dummy_renderer

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

       return components



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



 def       lowerCAmelCase_  (							self  :		List[Any]						):
       SCREAMING_SNAKE_CASE_							=			'cpu'

       SCREAMING_SNAKE_CASE_							=			self.get_dummy_components()

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

       pipe.set_progress_bar_config(disable=_lowerCAmelCase						)

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

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

       SCREAMING_SNAKE_CASE_							=			np.array(
           [
               0.0003_9216,
               0.0003_9216,
               0.0003_9216,
               0.0003_9216,
               0.0003_9216,
               0.0003_9216,
               0.0003_9216,
               0.0003_9216,
               0.0003_9216,
           ]						)

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



 def       lowerCAmelCase_  (							self  :		List[str]						):
       # NOTE: Larger batch sizes cause this test to timeout, only test on smaller batches
       self._test_inference_batch_consistent(batch_sizes=[1, 2]						)



 def       lowerCAmelCase_  (							self  :		int						):
       SCREAMING_SNAKE_CASE_							=			torch_device == 'cpu'
       SCREAMING_SNAKE_CASE_							=			True

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




 def       lowerCAmelCase_  (							self  :		List[Any]						):
       SCREAMING_SNAKE_CASE_							=			self.get_dummy_components()
       SCREAMING_SNAKE_CASE_							=			self.pipeline_class(**_lowerCAmelCase						)
       SCREAMING_SNAKE_CASE_							=			pipe.to(_lowerCAmelCase						)
       pipe.set_progress_bar_config(disable=_lowerCAmelCase						)

       SCREAMING_SNAKE_CASE_							=			1
       SCREAMING_SNAKE_CASE_							=			2

       SCREAMING_SNAKE_CASE_							=			self.get_dummy_inputs(_lowerCAmelCase						)

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

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

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





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


 '''simple docstring'''







 def       lowerCAmelCase_  (							self  :		Dict						):
       # clean up the VRAM after each test
       super().tearDown()
       gc.collect()
       torch.cuda.empty_cache()




 def       lowerCAmelCase_  (							self  :		Optional[Any]						):
       SCREAMING_SNAKE_CASE_							=			load_numpy(
           'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
           '/shap_e/test_shap_e_np_out.npy'						)
       SCREAMING_SNAKE_CASE_							=			ShapEPipeline.from_pretrained('openai/shap-e'						)
       SCREAMING_SNAKE_CASE_							=			pipe.to(_lowerCAmelCase						)
       pipe.set_progress_bar_config(disable=_lowerCAmelCase						)

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

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





import re

from filelock import FileLock


try:
   import nltk

   lowerCamelCase__ :   str					=							True
except (ImportError, ModuleNotFoundError):
   lowerCamelCase__ :   Union[str, Any]					=							False

if NLTK_AVAILABLE:
   with FileLock('.lock') as lock:
      nltk.download('punkt', quiet=True)

def 					UpperCAmelCase_							(      __UpperCAmelCase						:	str							)					-> str:
      re.sub('<n>'			,			''			,			__UpperCAmelCase							)  # remove pegasus newline char
      assert NLTK_AVAILABLE, "nltk must be installed to separate newlines between sentences. (pip install nltk)"
      return "\n".join(nltk.sent_tokenize(__UpperCAmelCase							)							)