"""simple docstring"""
import gc
import random
import unittest

import numpy as np
import torch
from transformers import CLIPImageProcessor, CLIPVisionConfig, CLIPVisionModel

from diffusers import HeunDiscreteScheduler, PriorTransformer, ShapEImgaImgPipeline
from diffusers.pipelines.shap_e import ShapERenderer
from diffusers.utils import floats_tensor, load_image, 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        __snake_case	(					_a	,				unittest.TestCase    ):
      a__				=		ShapEImgaImgPipeline
      a__				=		["""image"""]
      a__				=		["""image"""]
      a__				=		[
          """num_images_per_prompt""",
          """num_inference_steps""",
          """generator""",
          """latents""",
          """guidance_scale""",
          """frame_size""",
          """output_type""",
          """return_dict""",
      ]
      a__				=		False



      @property
      def   lowerCamelCase_       (						self)     ->      Optional[Any]:





        '''simple docstring'''





        return 32



      @property
      def   lowerCamelCase_       (						self)     ->      int:





        '''simple docstring'''





        return 32



      @property
      def   lowerCamelCase_       (						self)     ->      str:





        '''simple docstring'''





        return self.time_input_dim * 4



      @property
      def   lowerCamelCase_       (						self)     ->      List[Any]:





        '''simple docstring'''





        return 8



      @property
      def   lowerCamelCase_       (						self)     ->      Tuple:





        '''simple docstring'''





        torch.manual_seed(0)
        a__:					Tuple      							=     CLIPVisionConfig(
            hidden_size=self.text_embedder_hidden_size   ,      image_size=64   ,      projection_dim=self.text_embedder_hidden_size   ,      intermediate_size=37   ,      num_attention_heads=4   ,      num_channels=3   ,      num_hidden_layers=5   ,      patch_size=1   ,      )

        a__:					List[Any]      							=     CLIPVisionModel(__lowerCamelCase)
        return model



      @property
      def   lowerCamelCase_       (						self)     ->      Any:





        '''simple docstring'''





        a__:					Optional[int]      							=     CLIPImageProcessor(
            crop_size=2_24   ,      do_center_crop=__lowerCamelCase   ,      do_normalize=__lowerCamelCase   ,      do_resize=__lowerCamelCase   ,      image_mean=[0.48145466, 0.4578275, 0.40821073]   ,      image_std=[0.26862954, 0.26130258, 0.27577711]   ,      resample=3   ,      size=2_24   ,      )

        return image_processor



      @property
      def   lowerCamelCase_       (						self)     ->      Dict:





        '''simple docstring'''





        torch.manual_seed(0)

        a__:					Optional[int]      							=     {
            """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""",
            """embedding_proj_norm_type""": """layer""",
            """encoder_hid_proj_type""": None,
            """added_emb_type""": None,
        }

        a__:					int      							=     PriorTransformer(**__lowerCamelCase)
        return model



      @property
      def   lowerCamelCase_       (						self)     ->      int:





        '''simple docstring'''





        torch.manual_seed(0)

        a__:					List[str]      							=     {
            """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,
            ),
        }
        a__:					Tuple      							=     ShapERenderer(**__lowerCamelCase)
        return model



      def   lowerCamelCase_       (						self)     ->      Union[str, Any]:





        '''simple docstring'''





        a__:					Tuple      							=     self.dummy_prior
        a__:					Dict      							=     self.dummy_image_encoder
        a__:					List[str]      							=     self.dummy_image_processor
        a__:					List[Any]      							=     self.dummy_renderer

        a__:					Dict      							=     HeunDiscreteScheduler(
            beta_schedule='exp'   ,      num_train_timesteps=10_24   ,      prediction_type='sample'   ,      use_karras_sigmas=__lowerCamelCase   ,      clip_sample=__lowerCamelCase   ,      clip_sample_range=1.0   ,      )
        a__:					Tuple      							=     {
            """prior""": prior,
            """image_encoder""": image_encoder,
            """image_processor""": image_processor,
            """renderer""": renderer,
            """scheduler""": scheduler,
        }

        return components



      def   lowerCamelCase_       (						self   ,      lowercase   ,      lowercase=0)     ->      Optional[Any]:





        '''simple docstring'''





        a__:					Dict      							=     floats_tensor((1, 3, 64, 64)   ,      rng=random.Random(__lowerCamelCase)).to(__lowerCamelCase)

        if str(__lowerCamelCase).startswith('mps'):
          a__:					Optional[int]      							=     torch.manual_seed(__lowerCamelCase)
        else:
          a__:					Union[str, Any]      							=     torch.Generator(device=__lowerCamelCase).manual_seed(__lowerCamelCase)
        a__:					Tuple      							=     {
            """image""": input_image,
            """generator""": generator,
            """num_inference_steps""": 1,
            """frame_size""": 32,
            """output_type""": """np""",
        }
        return inputs



      def   lowerCamelCase_       (						self)     ->      int:





        '''simple docstring'''





        a__:					Dict      							=     """cpu"""

        a__:					List[str]      							=     self.get_dummy_components()

        a__:					Optional[int]      							=     self.pipeline_class(**__lowerCamelCase)
        a__:					Optional[int]      							=     pipe.to(__lowerCamelCase)

        pipe.set_progress_bar_config(disable=__lowerCamelCase)

        a__:					Union[str, Any]      							=     pipe(**self.get_dummy_inputs(__lowerCamelCase))
        a__:					Any      							=     output.images[0]
        a__:					str      							=     image[0, -3:, -3:, -1]

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

        a__:					List[str]      							=     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   lowerCamelCase_       (						self)     ->      Any:





        '''simple docstring'''





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



      def   lowerCamelCase_       (						self)     ->      Dict:





        '''simple docstring'''





        a__:					Union[str, Any]      							=     torch_device == """cpu"""
        a__:					Union[str, Any]      							=     True
        self._test_inference_batch_single_identical(
            batch_size=2   ,      test_max_difference=__lowerCamelCase   ,      relax_max_difference=__lowerCamelCase   ,      )



      def   lowerCamelCase_       (						self)     ->      int:





        '''simple docstring'''





        a__:					int      							=     self.get_dummy_components()
        a__:					List[str]      							=     self.pipeline_class(**__lowerCamelCase)
        a__:					Optional[Any]      							=     pipe.to(__lowerCamelCase)
        pipe.set_progress_bar_config(disable=__lowerCamelCase)

        a__:					Union[str, Any]      							=     1
        a__:					int      							=     2

        a__:					Optional[Any]      							=     self.get_dummy_inputs(__lowerCamelCase)

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

        a__:					Tuple      							=     pipe(**__lowerCamelCase   ,      num_images_per_prompt=__lowerCamelCase)[0]

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







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



      def   lowerCamelCase_       (						self)     ->      Optional[Any]:





        '''simple docstring'''





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



      def   lowerCamelCase_       (						self)     ->      Union[str, Any]:





        '''simple docstring'''





        a__:					Tuple      							=     load_image(
            'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main' '/shap_e/corgi.png')
        a__:					Tuple      							=     load_numpy(
            'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
            '/shap_e/test_shap_e_img2img_out.npy')
        a__:					Any      							=     ShapEImgaImgPipeline.from_pretrained('openai/shap-e-img2img')
        a__:					str      							=     pipe.to(__lowerCamelCase)
        pipe.set_progress_bar_config(disable=__lowerCamelCase)

        a__:					Any      							=     torch.Generator(device=__lowerCamelCase).manual_seed(0)

        a__:					Tuple      							=     pipe(
            __lowerCamelCase   ,      generator=__lowerCamelCase   ,      guidance_scale=3.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)



"""simple docstring"""
from __future__ import annotations

import unittest

import numpy as np

from transformers import BlipTextConfig
from transformers.testing_utils import require_tf, slow
from transformers.utils import is_tf_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor, random_attention_mask


if is_tf_available():
 import tensorflow as tf

 from transformers import TFBlipTextModel
 from transformers.models.blip.modeling_tf_blip import TF_BLIP_PRETRAINED_MODEL_ARCHIVE_LIST




class        __snake_case	:



      def __init__(						self   ,      lowercase   ,      lowercase=12   ,      lowercase=7   ,      lowercase=True   ,      lowercase=True   ,      lowercase=True   ,      lowercase=99   ,      lowercase=32   ,      lowercase=32   ,      lowercase=2   ,      lowercase=4   ,      lowercase=37   ,      lowercase=0.1   ,      lowercase=0.1   ,      lowercase=5_12   ,      lowercase=0.02   ,      lowercase=0   ,      lowercase=None   ,      )     ->      Optional[int]:





        '''simple docstring'''





        a__:					List[Any]      							=     parent
        a__:					Any      							=     batch_size
        a__:					int      							=     seq_length
        a__:					List[str]      							=     is_training
        a__:					Any      							=     use_input_mask
        a__:					Optional[Any]      							=     use_labels
        a__:					List[Any]      							=     vocab_size
        a__:					Optional[Any]      							=     hidden_size
        a__:					Any      							=     projection_dim
        a__:					List[str]      							=     num_hidden_layers
        a__:					Dict      							=     num_attention_heads
        a__:					int      							=     intermediate_size
        a__:					Tuple      							=     dropout
        a__:					Union[str, Any]      							=     attention_dropout
        a__:					str      							=     max_position_embeddings
        a__:					List[str]      							=     initializer_range
        a__:					Optional[int]      							=     scope
        a__:					List[Any]      							=     bos_token_id



      def   lowerCamelCase_       (						self)     ->      Dict:





        '''simple docstring'''





        a__:					str      							=     ids_tensor([self.batch_size, self.seq_length]   ,      self.vocab_size)

        a__:					str      							=     None
        if self.use_input_mask:
          a__:					List[str]      							=     random_attention_mask([self.batch_size, self.seq_length])

        if input_mask is not None:
          a__:					str      							=     input_mask.numpy()
          a__							,   a__:					Any      							=     input_mask.shape
          a__:					Tuple      							=     np.random.randint(1   ,      seq_length - 1   ,      size=(batch_size,))
          for batch_idx, start_index in enumerate(lowercase):
            a__:					Tuple      							=     1
            a__:					Any      							=     0

        a__:					Any      							=     self.get_config()

        return config, input_ids, tf.convert_to_tensor(lowercase)



      def   lowerCamelCase_       (						self)     ->      Union[str, Any]:





        '''simple docstring'''





        return BlipTextConfig(
            vocab_size=self.vocab_size   ,      hidden_size=self.hidden_size   ,      projection_dim=self.projection_dim   ,      num_hidden_layers=self.num_hidden_layers   ,      num_attention_heads=self.num_attention_heads   ,      intermediate_size=self.intermediate_size   ,      dropout=self.dropout   ,      attention_dropout=self.attention_dropout   ,      max_position_embeddings=self.max_position_embeddings   ,      initializer_range=self.initializer_range   ,      bos_token_id=self.bos_token_id   ,      )



      def   lowerCamelCase_       (						self   ,      lowercase   ,      lowercase   ,      lowercase)     ->      List[Any]:





        '''simple docstring'''





        a__:					Any      							=     TFBlipTextModel(config=lowercase)
        a__:					Union[str, Any]      							=     model(lowercase   ,      attention_mask=lowercase   ,      training=lowercase)
        a__:					List[Any]      							=     model(lowercase   ,      training=lowercase)
        self.parent.assertEqual(result.last_hidden_state.shape   ,      (self.batch_size, self.seq_length, self.hidden_size))
        self.parent.assertEqual(result.pooler_output.shape   ,      (self.batch_size, self.hidden_size))



      def   lowerCamelCase_       (						self)     ->      int:





        '''simple docstring'''





        a__:					Tuple      							=     self.prepare_config_and_inputs()
        a__							,   a__							,   a__:					Optional[Any]      							=     config_and_inputs
        a__:					Dict      							=     {'input_ids': input_ids, 'attention_mask': input_mask}
        return config, inputs_dict







@require_tf
class        __snake_case	(					__lowerCAmelCase	,				unittest.TestCase    ):
      a__				=		(TFBlipTextModel,) if is_tf_available() else ()
      a__				=		False
      a__				=		False
      a__				=		False



      def   lowerCamelCase_       (						self)     ->      Any:





        '''simple docstring'''





        a__:					List[Any]      							=     BlipTextModelTester(self)
        a__:					int      							=     ConfigTester(self   ,      config_class=lowercase   ,      hidden_size=37)



      def   lowerCamelCase_       (						self)     ->      Optional[Any]:





        '''simple docstring'''





        self.config_tester.run_common_tests()



      def   lowerCamelCase_       (						self)     ->      Union[str, Any]:





        '''simple docstring'''





        a__:					List[str]      							=     self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_model(*lowercase)



      def   lowerCamelCase_       (						self)     ->      Dict:





        '''simple docstring'''





        pass



      def   lowerCamelCase_       (						self)     ->      List[str]:





        '''simple docstring'''





        pass



      @unittest.skip(reason='Blip does not use inputs_embeds')
      def   lowerCamelCase_       (						self)     ->      List[str]:





        '''simple docstring'''





        pass



      @unittest.skip(reason='BlipTextModel has no base class and is not available in MODEL_MAPPING')
      def   lowerCamelCase_       (						self)     ->      Any:





        '''simple docstring'''





        pass



      @unittest.skip(reason='BlipTextModel has no base class and is not available in MODEL_MAPPING')
      def   lowerCamelCase_       (						self)     ->      Tuple:





        '''simple docstring'''





        pass



      @slow
      def   lowerCamelCase_       (						self)     ->      List[Any]:





        '''simple docstring'''





        for model_name in TF_BLIP_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
          a__:					Tuple      							=     TFBlipTextModel.from_pretrained(lowercase)
          self.assertIsNotNone(lowercase)



      def   lowerCamelCase_       (						self   ,      lowercase=True)     ->      str:





        '''simple docstring'''





        super().test_pt_tf_model_equivalence(allow_missing_keys=lowercase)


