





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 			lowerCAmelCase__( __lowercase  ,     unittest.TestCase    ):






       '''simple docstring'''


       __snake_case 			=    ShapEImgaImgPipeline
       __snake_case 			=    ['image']
       __snake_case 			=    ['image']
       __snake_case 			=    [
           'num_images_per_prompt',
           'num_inference_steps',
           'generator',
           'latents',
           'guidance_scale',
           'frame_size',
           'output_type',
           'return_dict',
       ]
       __snake_case 			=    False




       @property
       def     UpperCamelCase_			(      self							)			->     List[Any]:
            return 3_2




       @property
       def     UpperCamelCase_			(      self							)			->     List[str]:
            return 3_2




       @property
       def     UpperCamelCase_			(      self							)			->     str:
            return self.time_input_dim * 4




       @property
       def     UpperCamelCase_			(      self							)			->     List[Any]:
            return 8




       @property
       def     UpperCamelCase_			(      self							)			->     List[str]:
            torch.manual_seed(0							)
            _SCREAMING_SNAKE_CASE       :	Tuple						    =							CLIPVisionConfig(
                hidden_size=self.text_embedder_hidden_size    ,						image_size=6_4    ,						projection_dim=self.text_embedder_hidden_size    ,						intermediate_size=3_7    ,						num_attention_heads=4    ,						num_channels=3    ,						num_hidden_layers=5    ,						patch_size=1    ,						)

            _SCREAMING_SNAKE_CASE       :	Tuple						    =							CLIPVisionModel(__lowerCamelCase							)
            return model




       @property
       def     UpperCamelCase_			(      self							)			->     int:
            _SCREAMING_SNAKE_CASE       :	Union[str, Any]						    =							CLIPImageProcessor(
                crop_size=2_2_4    ,						do_center_crop=__lowerCamelCase    ,						do_normalize=__lowerCamelCase    ,						do_resize=__lowerCamelCase    ,						image_mean=[0.4814_5466, 0.457_8275, 0.4082_1073]    ,						image_std=[0.2686_2954, 0.2613_0258, 0.2757_7711]    ,						resample=3    ,						size=2_2_4    ,						)

            return image_processor




       @property
       def     UpperCamelCase_			(      self							)			->     List[str]:
            torch.manual_seed(0							)

            _SCREAMING_SNAKE_CASE       :	Union[str, Any]						    =							{
                "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",
                "embedding_proj_norm_type": "layer",
                "encoder_hid_proj_type": None,
                "added_emb_type": None,
            }

            _SCREAMING_SNAKE_CASE       :	List[str]						    =							PriorTransformer(**__lowerCamelCase							)
            return model




       @property
       def     UpperCamelCase_			(      self							)			->     Tuple:
            torch.manual_seed(0							)

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




       def     UpperCamelCase_			(      self							)			->     str:
            _SCREAMING_SNAKE_CASE       :	Optional[int]						    =							self.dummy_prior
            _SCREAMING_SNAKE_CASE       :	Any						    =							self.dummy_image_encoder
            _SCREAMING_SNAKE_CASE       :	List[str]						    =							self.dummy_image_processor
            _SCREAMING_SNAKE_CASE       :	Tuple						    =							self.dummy_renderer

            _SCREAMING_SNAKE_CASE       :	Union[str, Any]						    =							HeunDiscreteScheduler(
                beta_schedule="exp"    ,						num_train_timesteps=1_0_2_4    ,						prediction_type="sample"    ,						use_karras_sigmas=__lowerCamelCase    ,						clip_sample=__lowerCamelCase    ,						clip_sample_range=1.0    ,						)
            _SCREAMING_SNAKE_CASE       :	str						    =							{
                "prior": prior,
                "image_encoder": image_encoder,
                "image_processor": image_processor,
                "renderer": renderer,
                "scheduler": scheduler,
            }

            return components




       def     UpperCamelCase_			(      self    ,						__lowerCamelCase    ,						__lowerCamelCase=0							)			->     Optional[Any]:
            _SCREAMING_SNAKE_CASE       :	Union[str, Any]						    =							floats_tensor((1, 3, 6_4, 6_4)    ,						rng=random.Random(__lowerCamelCase							)							).to(__lowerCamelCase							)

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




       def     UpperCamelCase_			(      self							)			->     Optional[Any]:
            _SCREAMING_SNAKE_CASE       :	int						    =							"cpu"

            _SCREAMING_SNAKE_CASE       :	Dict						    =							self.get_dummy_components()

            _SCREAMING_SNAKE_CASE       :	Any						    =							self.pipeline_class(**__lowerCamelCase							)
            _SCREAMING_SNAKE_CASE       :	List[Any]						    =							pipe.to(__lowerCamelCase							)

            pipe.set_progress_bar_config(disable=__lowerCamelCase							)

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

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

            _SCREAMING_SNAKE_CASE       :	Union[str, Any]						    =							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     UpperCamelCase_			(      self							)			->     Optional[int]:
            # NOTE: Larger batch sizes cause this test to timeout, only test on smaller batches
            self._test_inference_batch_consistent(batch_sizes=[1, 2]							)




       def     UpperCamelCase_			(      self							)			->     Dict:
            _SCREAMING_SNAKE_CASE       :	int						    =							torch_device == "cpu"
            _SCREAMING_SNAKE_CASE       :	Dict						    =							True
            self._test_inference_batch_single_identical(
                batch_size=2    ,						test_max_difference=__lowerCamelCase    ,						relax_max_difference=__lowerCamelCase    ,						)







       def     UpperCamelCase_			(      self							)			->     Optional[Any]:
            _SCREAMING_SNAKE_CASE       :	Union[str, Any]						    =							self.get_dummy_components()
            _SCREAMING_SNAKE_CASE       :	List[str]						    =							self.pipeline_class(**__lowerCamelCase							)
            _SCREAMING_SNAKE_CASE       :	Optional[int]						    =							pipe.to(__lowerCamelCase							)
            pipe.set_progress_bar_config(disable=__lowerCamelCase							)

            _SCREAMING_SNAKE_CASE       :	Dict						    =							1
            _SCREAMING_SNAKE_CASE       :	List[Any]						    =							2

            _SCREAMING_SNAKE_CASE       :	List[str]						    =							self.get_dummy_inputs(__lowerCamelCase							)

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

            _SCREAMING_SNAKE_CASE       :	Optional[int]						    =							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     UpperCamelCase_			(      self							)			->     str:
            # clean up the VRAM after each test
            super().tearDown()
            gc.collect()
            torch.cuda.empty_cache()







       def     UpperCamelCase_			(      self							)			->     Tuple:
            _SCREAMING_SNAKE_CASE       :	Optional[Any]						    =							load_image(
                "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main" "/shap_e/corgi.png"							)
            _SCREAMING_SNAKE_CASE       :	Tuple						    =							load_numpy(
                "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
                "/shap_e/test_shap_e_img2img_out.npy"							)
            _SCREAMING_SNAKE_CASE       :	int						    =							ShapEImgaImgPipeline.from_pretrained("openai/shap-e-img2img"							)
            _SCREAMING_SNAKE_CASE       :	str						    =							pipe.to(__lowerCamelCase							)
            pipe.set_progress_bar_config(disable=__lowerCamelCase							)

            _SCREAMING_SNAKE_CASE       :	Optional[Any]						    =							torch.Generator(device=__lowerCamelCase							).manual_seed(0							)

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






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 			lowerCAmelCase__:






       '''simple docstring'''


       def __init__(      self    ,						__lowerCamelCase    ,						__lowerCamelCase=1_2    ,						__lowerCamelCase=7    ,						__lowerCamelCase=True    ,						__lowerCamelCase=True    ,						__lowerCamelCase=True    ,						__lowerCamelCase=9_9    ,						__lowerCamelCase=3_2    ,						__lowerCamelCase=3_2    ,						__lowerCamelCase=2    ,						__lowerCamelCase=4    ,						__lowerCamelCase=3_7    ,						__lowerCamelCase=0.1    ,						__lowerCamelCase=0.1    ,						__lowerCamelCase=5_1_2    ,						__lowerCamelCase=0.02    ,						__lowerCamelCase=0    ,						__lowerCamelCase=None    ,						)			->     Any:
            _SCREAMING_SNAKE_CASE       :	int						    =							parent
            _SCREAMING_SNAKE_CASE       :	Union[str, Any]						    =							batch_size
            _SCREAMING_SNAKE_CASE       :	str						    =							seq_length
            _SCREAMING_SNAKE_CASE       :	str						    =							is_training
            _SCREAMING_SNAKE_CASE       :	str						    =							use_input_mask
            _SCREAMING_SNAKE_CASE       :	int						    =							use_labels
            _SCREAMING_SNAKE_CASE       :	List[Any]						    =							vocab_size
            _SCREAMING_SNAKE_CASE       :	Dict						    =							hidden_size
            _SCREAMING_SNAKE_CASE       :	int						    =							projection_dim
            _SCREAMING_SNAKE_CASE       :	Optional[Any]						    =							num_hidden_layers
            _SCREAMING_SNAKE_CASE       :	int						    =							num_attention_heads
            _SCREAMING_SNAKE_CASE       :	Optional[int]						    =							intermediate_size
            _SCREAMING_SNAKE_CASE       :	int						    =							dropout
            _SCREAMING_SNAKE_CASE       :	Union[str, Any]						    =							attention_dropout
            _SCREAMING_SNAKE_CASE       :	Optional[int]						    =							max_position_embeddings
            _SCREAMING_SNAKE_CASE       :	Dict						    =							initializer_range
            _SCREAMING_SNAKE_CASE       :	Tuple						    =							scope
            _SCREAMING_SNAKE_CASE       :	Any						    =							bos_token_id




       def     UpperCamelCase_			(      self							)			->     Optional[Any]:
            _SCREAMING_SNAKE_CASE       :	Any						    =							ids_tensor([self.batch_size, self.seq_length]    ,						self.vocab_size							)

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

            if input_mask is not None:
                 _SCREAMING_SNAKE_CASE       :	Optional[Any]						    =							input_mask.numpy()
                 _SCREAMING_SNAKE_CASE    ,			_SCREAMING_SNAKE_CASE       :	int						    =							input_mask.shape
                 _SCREAMING_SNAKE_CASE       :	str						    =							np.random.randint(1    ,						seq_length - 1    ,						size=(batch_size,)							)
                 for batch_idx, start_index in enumerate(__lowerCamelCase							):
                      _SCREAMING_SNAKE_CASE       :	Union[str, Any]						    =							1
                      _SCREAMING_SNAKE_CASE       :	List[str]						    =							0

            _SCREAMING_SNAKE_CASE       :	Optional[int]						    =							self.get_config()

            return config, input_ids, tf.convert_to_tensor(__lowerCamelCase							)




       def     UpperCamelCase_			(      self							)			->     List[str]:
            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     UpperCamelCase_			(      self    ,						__lowerCamelCase    ,						__lowerCamelCase    ,						__lowerCamelCase							)			->     List[str]:
            _SCREAMING_SNAKE_CASE       :	Optional[int]						    =							TFBlipTextModel(config=__lowerCamelCase							)
            _SCREAMING_SNAKE_CASE       :	Union[str, Any]						    =							model(__lowerCamelCase    ,						attention_mask=__lowerCamelCase    ,						training=__lowerCamelCase							)
            _SCREAMING_SNAKE_CASE       :	Dict						    =							model(__lowerCamelCase    ,						training=__lowerCamelCase							)
            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     UpperCamelCase_			(      self							)			->     str:
            _SCREAMING_SNAKE_CASE       :	Optional[int]						    =							self.prepare_config_and_inputs()
            _SCREAMING_SNAKE_CASE    ,			_SCREAMING_SNAKE_CASE    ,			_SCREAMING_SNAKE_CASE       :	int						    =							config_and_inputs
            _SCREAMING_SNAKE_CASE       :	List[str]						    =							{"input_ids": input_ids, "attention_mask": input_mask}
            return config, inputs_dict





@require_tf
class 			lowerCAmelCase__( __lowercase  ,     unittest.TestCase    ):






       '''simple docstring'''


       __snake_case 			=    (TFBlipTextModel,) if is_tf_available() else ()
       __snake_case 			=    False
       __snake_case 			=    False
       __snake_case 			=    False




       def     UpperCamelCase_			(      self							)			->     Any:
            _SCREAMING_SNAKE_CASE       :	Optional[Any]						    =							BlipTextModelTester(self							)
            _SCREAMING_SNAKE_CASE       :	Tuple						    =							ConfigTester(self    ,						config_class=__lowerCamelCase    ,						hidden_size=3_7							)




       def     UpperCamelCase_			(      self							)			->     Union[str, Any]:
            self.config_tester.run_common_tests()




       def     UpperCamelCase_			(      self							)			->     int:
            _SCREAMING_SNAKE_CASE       :	Optional[int]						    =							self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_model(*__lowerCamelCase							)




       def     UpperCamelCase_			(      self							)			->     int:
            pass




       def     UpperCamelCase_			(      self							)			->     Tuple:
            pass




       @unittest.skip(reason="Blip does not use inputs_embeds"							)
       def     UpperCamelCase_			(      self							)			->     Dict:
            pass




       @unittest.skip(reason="BlipTextModel has no base class and is not available in MODEL_MAPPING"							)
       def     UpperCamelCase_			(      self							)			->     Optional[int]:
            pass




       @unittest.skip(reason="BlipTextModel has no base class and is not available in MODEL_MAPPING"							)
       def     UpperCamelCase_			(      self							)			->     Optional[Any]:
            pass




       @slow
       def     UpperCamelCase_			(      self							)			->     int:
            for model_name in TF_BLIP_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                 _SCREAMING_SNAKE_CASE       :	Union[str, Any]						    =							TFBlipTextModel.from_pretrained(__lowerCamelCase							)
                 self.assertIsNotNone(__lowerCamelCase							)







       def     UpperCamelCase_			(      self    ,						__lowerCamelCase=True							)			->     int:
            super().test_pt_tf_model_equivalence(allow_missing_keys=__lowerCamelCase							)