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




							@property
							def  A__      (				self):
												return 3_2




							@property
							def  A__      (				self):
												return 3_2




							@property
							def  A__      (				self):
												return self.time_input_dim * 4




							@property
							def  A__      (				self):
												return 8




							@property
							def  A__      (				self):
												lowercase							=						CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip''')
												return tokenizer




							@property
							def  A__      (				self):
												torch.manual_seed(0)
												lowercase							=						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  A__      (				self):
												torch.manual_seed(0)

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

												lowercase							=						PriorTransformer(**A__)
												return model




							@property
							def  A__      (				self):
												torch.manual_seed(0)

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




							def  A__      (				self):
												lowercase							=						self.dummy_prior
												lowercase							=						self.dummy_text_encoder
												lowercase							=						self.dummy_tokenizer
												lowercase							=						self.dummy_renderer

												lowercase							=						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   ,)
												lowercase							=						{
												    '''prior''': prior,
												    '''text_encoder''': text_encoder,
												    '''tokenizer''': tokenizer,
												    '''renderer''': renderer,
												    '''scheduler''': scheduler,
												}

												return components




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




							def  A__      (				self):
												lowercase							=						'''cpu'''

												lowercase							=						self.get_dummy_components()

												lowercase							=						self.pipeline_class(**A__)
												lowercase							=						pipe.to(A__)

												pipe.set_progress_bar_config(disable=A__)

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

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

												lowercase							=						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  A__      (				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  A__      (				self):
												lowercase							=						torch_device == '''cpu'''
												lowercase							=						True

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





							def  A__      (				self):
												lowercase							=						self.get_dummy_components()
												lowercase							=						self.pipeline_class(**A__)
												lowercase							=						pipe.to(A__)
												pipe.set_progress_bar_config(disable=A__)

												lowercase							=						1
												lowercase							=						2

												lowercase							=						self.get_dummy_inputs(A__)

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

												lowercase							=						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  A__      (				self):
												# clean up the VRAM after each test
												super().tearDown()
												gc.collect()
												torch.cuda.empty_cache()





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

												lowercase							=						torch.Generator(device=A__).manual_seed(0)

												lowercase							=						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__)





"""simple docstring"""







import enum
import warnings

from ..tokenization_utils import TruncationStrategy
from ..utils import add_end_docstrings, is_tf_available, is_torch_available, logging
from .base import PIPELINE_INIT_ARGS, Pipeline


if is_tf_available():
       import tensorflow as tf

       from ..models.auto.modeling_tf_auto import TF_MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING

if is_torch_available():
       from ..models.auto.modeling_auto import MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING

__SCREAMING_SNAKE_CASE    	=logging.get_logger(__name__)







class        UpperCamelCase		(			enum.Enum  ):
   lowercase							=  0
   lowercase							=  1







@add_end_docstrings(lowercase_  )
class        UpperCamelCase		(			lowercase_  ):
   lowercase							=  'generated'
   def __init__(			self     ,*__UpperCamelCase     ,**__UpperCamelCase			)	->      Tuple:


        '''simple docstring'''


        super().__init__(*__UpperCamelCase     ,**__UpperCamelCase			)

        self.check_model_type(
            TF_MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING
            if self.framework == 'tf'
            else MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING			)
   def 			_UpperCAmelCase   (			self     ,__UpperCamelCase=None     ,__UpperCamelCase=None     ,__UpperCamelCase=None     ,__UpperCamelCase=None     ,__UpperCamelCase=None     ,__UpperCamelCase=None     ,**__UpperCamelCase     ,)	->      Optional[Any]:


        '''simple docstring'''


        lowercase_	:				List[Any]													=	{}
        if truncation is not None:
             lowercase_	:				int													=	truncation

        lowercase_	:				Dict													=	generate_kwargs

        lowercase_	:				List[Any]													=	{}
        if return_tensors is not None and return_type is None:
             lowercase_	:				Union[str, Any]													=	ReturnType.TENSORS if return_tensors else ReturnType.TEXT
        if return_type is not None:
             lowercase_	:				str													=	return_type

        if clean_up_tokenization_spaces is not None:
             lowercase_	:				Dict													=	clean_up_tokenization_spaces

        if stop_sequence is not None:
             lowercase_	:				Union[str, Any]													=	self.tokenizer.encode(__UpperCamelCase     ,add_special_tokens=__UpperCamelCase			)
             if len(__UpperCamelCase			) > 1:
                  warnings.warn(
                      'Stopping on a multiple token sequence is not yet supported on transformers. The first token of'
                      ' the stop sequence will be used as the stop sequence string in the interim.'			)
             lowercase_	:				Optional[int]													=	stop_sequence_ids[0]

        return preprocess_params, forward_params, postprocess_params
   def 			_UpperCAmelCase   (			self     ,__UpperCamelCase     ,__UpperCamelCase     ,__UpperCamelCase			)	->      Optional[int]:


        '''simple docstring'''


        return True
   def 			_UpperCAmelCase   (			self     ,*__UpperCamelCase     ,__UpperCamelCase			)	->      Dict:


        '''simple docstring'''


        lowercase_	:				Dict													=	self.model.config.prefix if self.model.config.prefix is not None else ''
        if isinstance(args[0]     ,__UpperCamelCase			):
             if self.tokenizer.pad_token_id is None:
                  raise ValueError('Please make sure that the tokenizer has a pad_token_id when using a batch input'			)
             lowercase_	:				str													=	([prefix + arg for arg in args[0]],)
             lowercase_	:				Union[str, Any]													=	True

        elif isinstance(args[0]     ,__UpperCamelCase			):
             lowercase_	:				Union[str, Any]													=	(prefix + args[0],)
             lowercase_	:				Union[str, Any]													=	False
        else:
             raise ValueError(
                 f''' `args[0]`: {args[0]} have the wrong format. The should be either of type `str` or type `list`'''			)
        lowercase_	:				List[Any]													=	self.tokenizer(*__UpperCamelCase     ,padding=__UpperCamelCase     ,truncation=__UpperCamelCase     ,return_tensors=self.framework			)
        # This is produced by tokenizers but is an invalid generate kwargs
        if "token_type_ids" in inputs:
             del inputs["token_type_ids"]
        return inputs
   def __call__(			self     ,*__UpperCamelCase     ,**__UpperCamelCase			)	->      Dict:


        '''simple docstring'''


        lowercase_	:				Optional[int]													=	super().__call__(*__UpperCamelCase     ,**__UpperCamelCase			)
        if (
            isinstance(args[0]     ,__UpperCamelCase			)
            and all(isinstance(__UpperCamelCase     ,__UpperCamelCase			) for el in args[0]			)
            and all(len(__UpperCamelCase			) == 1 for res in result			)
        ):
             return [res[0] for res in result]
        return result
   def 			_UpperCAmelCase   (			self     ,__UpperCamelCase     ,__UpperCamelCase=TruncationStrategy.DO_NOT_TRUNCATE     ,**__UpperCamelCase			)	->      List[str]:


        '''simple docstring'''


        lowercase_	:				Any													=	self._parse_and_tokenize(__UpperCamelCase     ,truncation=__UpperCamelCase     ,**__UpperCamelCase			)
        return inputs
   def 			_UpperCAmelCase   (			self     ,__UpperCamelCase     ,**__UpperCamelCase			)	->      Optional[int]:


        '''simple docstring'''


        if self.framework == "pt":
             lowercase_							,  lowercase_	:				Optional[int]													=	model_inputs['input_ids'].shape
        elif self.framework == "tf":
             lowercase_							,  lowercase_	:				Union[str, Any]													=	tf.shape(model_inputs['input_ids']			).numpy()

        lowercase_	:				str													=	generate_kwargs.get('min_length'     ,self.model.config.min_length			)
        lowercase_	:				List[Any]													=	generate_kwargs.get('max_length'     ,self.model.config.max_length			)
        self.check_inputs(__UpperCamelCase     ,generate_kwargs['min_length']     ,generate_kwargs['max_length']			)
        lowercase_	:				Tuple													=	self.model.generate(**__UpperCamelCase     ,**__UpperCamelCase			)
        lowercase_	:				str													=	output_ids.shape[0]
        if self.framework == "pt":
             lowercase_	:				List[Any]													=	output_ids.reshape(__UpperCamelCase     ,out_b // in_b     ,*output_ids.shape[1:]			)
        elif self.framework == "tf":
             lowercase_	:				List[Any]													=	tf.reshape(__UpperCamelCase     ,(in_b, out_b // in_b, *output_ids.shape[1:])			)
        return {"output_ids": output_ids}







   def 			_UpperCAmelCase   (			self     ,__UpperCamelCase     ,__UpperCamelCase=ReturnType.TEXT     ,__UpperCamelCase=False			)	->      Dict:


        '''simple docstring'''


        lowercase_	:				Dict													=	[]
        for output_ids in model_outputs["output_ids"][0]:
             if return_type == ReturnType.TENSORS:
                  lowercase_	:				List[Any]													=	{f'''{self.return_name}_token_ids''': output_ids}
             elif return_type == ReturnType.TEXT:
                  lowercase_	:				str													=	{
                      f'''{self.return_name}_text''': self.tokenizer.decode(
                          __UpperCamelCase     ,skip_special_tokens=__UpperCamelCase     ,clean_up_tokenization_spaces=__UpperCamelCase     ,)
                  }
             records.append(__UpperCamelCase			)
        return records







@add_end_docstrings(lowercase_  )
class        UpperCamelCase		(			lowercase_  ):
   lowercase							=  'summary'
   def __call__(			self     ,*__UpperCamelCase     ,**__UpperCamelCase			)	->      Union[str, Any]:


        '''simple docstring'''


        return super().__call__(*__UpperCamelCase     ,**__UpperCamelCase			)







   def 			_UpperCAmelCase   (			self     ,__UpperCamelCase     ,__UpperCamelCase     ,__UpperCamelCase			)	->      bool:


        '''simple docstring'''


        if max_length < min_length:
             logger.warning(f'''Your min_length={min_length} must be inferior than your max_length={max_length}.'''			)

        if input_length < max_length:
             logger.warning(
                 f'''Your max_length is set to {max_length}, but your input_length is only {input_length}. Since this is '''
                 'a summarization task, where outputs shorter than the input are typically wanted, you might '
                 f'''consider decreasing max_length manually, e.g. summarizer(\'...\', max_length={input_length//2})'''			)







@add_end_docstrings(lowercase_  )
class        UpperCamelCase		(			lowercase_  ):
   lowercase							=  'translation'
   def 			_UpperCAmelCase   (			self     ,__UpperCamelCase     ,__UpperCamelCase     ,__UpperCamelCase			)	->      Dict:


        '''simple docstring'''


        if input_length > 0.9 * max_length:
             logger.warning(
                 f'''Your input_length: {input_length} is bigger than 0.9 * max_length: {max_length}. You might consider '''
                 'increasing your max_length manually, e.g. translator(\'...\', max_length=400)'			)
        return True
   def 			_UpperCAmelCase   (			self     ,*__UpperCamelCase     ,__UpperCamelCase=TruncationStrategy.DO_NOT_TRUNCATE     ,__UpperCamelCase=None     ,__UpperCamelCase=None			)	->      int:


        '''simple docstring'''


        if getattr(self.tokenizer     ,'_build_translation_inputs'     ,__UpperCamelCase			):
             return self.tokenizer._build_translation_inputs(
                 *__UpperCamelCase     ,return_tensors=self.framework     ,truncation=__UpperCamelCase     ,src_lang=__UpperCamelCase     ,tgt_lang=__UpperCamelCase			)
        else:
             return super()._parse_and_tokenize(*__UpperCamelCase     ,truncation=__UpperCamelCase			)
   def 			_UpperCAmelCase   (			self     ,__UpperCamelCase=None     ,__UpperCamelCase=None     ,**__UpperCamelCase			)	->      str:


        '''simple docstring'''


        lowercase_							,  lowercase_							,  lowercase_	:				int													=	super()._sanitize_parameters(**__UpperCamelCase			)
        if src_lang is not None:
             lowercase_	:				str													=	src_lang
        if tgt_lang is not None:
             lowercase_	:				Optional[Any]													=	tgt_lang
        if src_lang is None and tgt_lang is None:
             # Backward compatibility, direct arguments use is preferred.
             lowercase_	:				Tuple													=	kwargs.get('task'     ,self.task			)
             lowercase_	:				List[str]													=	task.split('_'			)
             if task and len(__UpperCamelCase			) == 4:
                  # translation, XX, to YY
                  lowercase_	:				Union[str, Any]													=	items[1]
                  lowercase_	:				Tuple													=	items[3]
        return preprocess_params, forward_params, postprocess_params







   def __call__(			self     ,*__UpperCamelCase     ,**__UpperCamelCase			)	->      Dict:


        '''simple docstring'''


        return super().__call__(*__UpperCamelCase     ,**__UpperCamelCase			)

