





"""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 		lowercase					:





      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=512      ,						lowercase=0.02      ,						lowercase=0      ,						lowercase=None      ,						)		->     Any:
            lowerCAmelCase  						=       parent
            lowerCAmelCase  						=       batch_size
            lowerCAmelCase  						=       seq_length
            lowerCAmelCase  						=       is_training
            lowerCAmelCase  						=       use_input_mask
            lowerCAmelCase  						=       use_labels
            lowerCAmelCase  						=       vocab_size
            lowerCAmelCase  						=       hidden_size
            lowerCAmelCase  						=       projection_dim
            lowerCAmelCase  						=       num_hidden_layers
            lowerCAmelCase  						=       num_attention_heads
            lowerCAmelCase  						=       intermediate_size
            lowerCAmelCase  						=       dropout
            lowerCAmelCase  						=       attention_dropout
            lowerCAmelCase  						=       max_position_embeddings
            lowerCAmelCase  						=       initializer_range
            lowerCAmelCase  						=       scope
            lowerCAmelCase  						=       bos_token_id





      def     _snake_case			(       self )		->     Optional[int]:
            lowerCAmelCase  						=       ids_tensor([self.batch_size, self.seq_length]      ,						self.vocab_size )

            lowerCAmelCase  						=       None
            if self.use_input_mask:
                  lowerCAmelCase  						=       random_attention_mask([self.batch_size, self.seq_length] )

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

            lowerCAmelCase  						=       self.get_config()

            return config, input_ids, tf.convert_to_tensor(lowercase )





      def     _snake_case			(       self )		->     Optional[int]:
            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     _snake_case			(       self      ,						lowercase      ,						lowercase      ,						lowercase )		->     Optional[int]:
            lowerCAmelCase  						=       TFBlipTextModel(config=lowercase )
            lowerCAmelCase  						=       model(lowercase      ,						attention_mask=lowercase      ,						training=lowercase )
            lowerCAmelCase  						=       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     _snake_case			(       self )		->     Tuple:
            lowerCAmelCase  						=       self.prepare_config_and_inputs()
            lowerCAmelCase						,		lowerCAmelCase						,		lowerCAmelCase  						=       config_and_inputs
            lowerCAmelCase  						=       {"""input_ids""": input_ids, """attention_mask""": input_mask}
            return config, inputs_dict







@require_tf
class 		lowercase					(					_UpperCAmelCase     ,  unittest.TestCase					):
      _SCREAMING_SNAKE_CASE					=      (TFBlipTextModel,) if is_tf_available() else ()
      _SCREAMING_SNAKE_CASE					=      False
      _SCREAMING_SNAKE_CASE					=      False
      _SCREAMING_SNAKE_CASE					=      False





      def     _snake_case			(       self )		->     List[Any]:
            lowerCAmelCase  						=       BlipTextModelTester(self )
            lowerCAmelCase  						=       ConfigTester(self      ,						config_class=lowercase      ,						hidden_size=37 )





      def     _snake_case			(       self )		->     Any:
            self.config_tester.run_common_tests()





      def     _snake_case			(       self )		->     str:
            lowerCAmelCase  						=       self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_model(*lowercase )





      def     _snake_case			(       self )		->     List[str]:
            pass





      def     _snake_case			(       self )		->     Tuple:
            pass





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





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





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





      @slow
      def     _snake_case			(       self )		->     List[Any]:
            for model_name in TF_BLIP_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                  lowerCAmelCase  						=       TFBlipTextModel.from_pretrained(lowercase )
                  self.assertIsNotNone(lowercase )





      def     _snake_case			(       self      ,						lowercase=True )		->     List[Any]:
            super().test_pt_tf_model_equivalence(allow_missing_keys=lowercase )













"""simple docstring"""





import tempfile

import numpy as np
import torch
from transformers import AutoTokenizer, TaEncoderModel

from diffusers import DDPMScheduler, UNetaDConditionModel
from diffusers.models.attention_processor import AttnAddedKVProcessor
from diffusers.pipelines.deepfloyd_if import IFWatermarker
from diffusers.utils.testing_utils import torch_device

from ..test_pipelines_common import to_np

class 		lowercase					:





      def     _snake_case			(       self )		->     List[Any]:
            torch.manual_seed(0 )
            lowerCAmelCase  						=       TaEncoderModel.from_pretrained("""hf-internal-testing/tiny-random-t5""" )

            torch.manual_seed(0 )
            lowerCAmelCase  						=       AutoTokenizer.from_pretrained("""hf-internal-testing/tiny-random-t5""" )

            torch.manual_seed(0 )
            lowerCAmelCase  						=       UNetaDConditionModel(
                sample_size=32      ,						layers_per_block=1      ,						block_out_channels=[32, 64]      ,						down_block_types=[
                    """ResnetDownsampleBlock2D""",
                    """SimpleCrossAttnDownBlock2D""",
                ]      ,						mid_block_type="""UNetMidBlock2DSimpleCrossAttn"""      ,						up_block_types=["""SimpleCrossAttnUpBlock2D""", """ResnetUpsampleBlock2D"""]      ,						in_channels=3      ,						out_channels=6      ,						cross_attention_dim=32      ,						encoder_hid_dim=32      ,						attention_head_dim=8      ,						addition_embed_type="""text"""      ,						addition_embed_type_num_heads=2      ,						cross_attention_norm="""group_norm"""      ,						resnet_time_scale_shift="""scale_shift"""      ,						act_fn="""gelu"""      ,						)
            unet.set_attn_processor(AttnAddedKVProcessor() )  # For reproducibility tests

            torch.manual_seed(0 )
            lowerCAmelCase  						=       DDPMScheduler(
                num_train_timesteps=1_000      ,						beta_schedule="""squaredcos_cap_v2"""      ,						beta_start=0.0_001      ,						beta_end=0.02      ,						thresholding=lowercase      ,						dynamic_thresholding_ratio=0.95      ,						sample_max_value=1.0      ,						prediction_type="""epsilon"""      ,						variance_type="""learned_range"""      ,						)

            torch.manual_seed(0 )
            lowerCAmelCase  						=       IFWatermarker()

            return {
                "text_encoder": text_encoder,
                "tokenizer": tokenizer,
                "unet": unet,
                "scheduler": scheduler,
                "watermarker": watermarker,
                "safety_checker": None,
                "feature_extractor": None,
            }





      def     _snake_case			(       self )		->     int:
            torch.manual_seed(0 )
            lowerCAmelCase  						=       TaEncoderModel.from_pretrained("""hf-internal-testing/tiny-random-t5""" )

            torch.manual_seed(0 )
            lowerCAmelCase  						=       AutoTokenizer.from_pretrained("""hf-internal-testing/tiny-random-t5""" )

            torch.manual_seed(0 )
            lowerCAmelCase  						=       UNetaDConditionModel(
                sample_size=32      ,						layers_per_block=[1, 2]      ,						block_out_channels=[32, 64]      ,						down_block_types=[
                    """ResnetDownsampleBlock2D""",
                    """SimpleCrossAttnDownBlock2D""",
                ]      ,						mid_block_type="""UNetMidBlock2DSimpleCrossAttn"""      ,						up_block_types=["""SimpleCrossAttnUpBlock2D""", """ResnetUpsampleBlock2D"""]      ,						in_channels=6      ,						out_channels=6      ,						cross_attention_dim=32      ,						encoder_hid_dim=32      ,						attention_head_dim=8      ,						addition_embed_type="""text"""      ,						addition_embed_type_num_heads=2      ,						cross_attention_norm="""group_norm"""      ,						resnet_time_scale_shift="""scale_shift"""      ,						act_fn="""gelu"""      ,						class_embed_type="""timestep"""      ,						mid_block_scale_factor=1.414      ,						time_embedding_act_fn="""gelu"""      ,						time_embedding_dim=32      ,						)
            unet.set_attn_processor(AttnAddedKVProcessor() )  # For reproducibility tests

            torch.manual_seed(0 )
            lowerCAmelCase  						=       DDPMScheduler(
                num_train_timesteps=1_000      ,						beta_schedule="""squaredcos_cap_v2"""      ,						beta_start=0.0_001      ,						beta_end=0.02      ,						thresholding=lowercase      ,						dynamic_thresholding_ratio=0.95      ,						sample_max_value=1.0      ,						prediction_type="""epsilon"""      ,						variance_type="""learned_range"""      ,						)

            torch.manual_seed(0 )
            lowerCAmelCase  						=       DDPMScheduler(
                num_train_timesteps=1_000      ,						beta_schedule="""squaredcos_cap_v2"""      ,						beta_start=0.0_001      ,						beta_end=0.02      ,						)

            torch.manual_seed(0 )
            lowerCAmelCase  						=       IFWatermarker()

            return {
                "text_encoder": text_encoder,
                "tokenizer": tokenizer,
                "unet": unet,
                "scheduler": scheduler,
                "image_noising_scheduler": image_noising_scheduler,
                "watermarker": watermarker,
                "safety_checker": None,
                "feature_extractor": None,
            }





      def     _snake_case			(       self )		->     Union[str, Any]:
            lowerCAmelCase  						=       self.get_dummy_components()
            lowerCAmelCase  						=       self.pipeline_class(**lowercase )
            pipe.to(lowercase )
            pipe.set_progress_bar_config(disable=lowercase )

            lowerCAmelCase  						=       self.get_dummy_inputs(lowercase )

            lowerCAmelCase  						=       inputs["""prompt"""]
            lowerCAmelCase  						=       inputs["""generator"""]
            lowerCAmelCase  						=       inputs["""num_inference_steps"""]
            lowerCAmelCase  						=       inputs["""output_type"""]

            if "image" in inputs:
                  lowerCAmelCase  						=       inputs["""image"""]
            else:
                  lowerCAmelCase  						=       None

            if "mask_image" in inputs:
                  lowerCAmelCase  						=       inputs["""mask_image"""]
            else:
                  lowerCAmelCase  						=       None

            if "original_image" in inputs:
                  lowerCAmelCase  						=       inputs["""original_image"""]
            else:
                  lowerCAmelCase  						=       None

            lowerCAmelCase						,		lowerCAmelCase  						=       pipe.encode_prompt(lowercase )

            # inputs with prompt converted to embeddings
            lowerCAmelCase  						=       {
                """prompt_embeds""": prompt_embeds,
                """negative_prompt_embeds""": negative_prompt_embeds,
                """generator""": generator,
                """num_inference_steps""": num_inference_steps,
                """output_type""": output_type,
            }

            if image is not None:
                  lowerCAmelCase  						=       image

            if mask_image is not None:
                  lowerCAmelCase  						=       mask_image

            if original_image is not None:
                  lowerCAmelCase  						=       original_image

            # set all optional components to None
            for optional_component in pipe._optional_components:
                  setattr(lowercase      ,						lowercase      ,						lowercase )

            lowerCAmelCase  						=       pipe(**lowercase )[0]

            with tempfile.TemporaryDirectory() as tmpdir:
                  pipe.save_pretrained(lowercase )
                  lowerCAmelCase  						=       self.pipeline_class.from_pretrained(lowercase )
                  pipe_loaded.to(lowercase )
                  pipe_loaded.set_progress_bar_config(disable=lowercase )

            pipe_loaded.unet.set_attn_processor(AttnAddedKVProcessor() )  # For reproducibility tests

            for optional_component in pipe._optional_components:
                  self.assertTrue(
                      getattr(lowercase      ,						lowercase ) is None      ,						f'`{optional_component}` did not stay set to None after loading.'      ,						)

            lowerCAmelCase  						=       self.get_dummy_inputs(lowercase )

            lowerCAmelCase  						=       inputs["""generator"""]
            lowerCAmelCase  						=       inputs["""num_inference_steps"""]
            lowerCAmelCase  						=       inputs["""output_type"""]

            # inputs with prompt converted to embeddings
            lowerCAmelCase  						=       {
                """prompt_embeds""": prompt_embeds,
                """negative_prompt_embeds""": negative_prompt_embeds,
                """generator""": generator,
                """num_inference_steps""": num_inference_steps,
                """output_type""": output_type,
            }

            if image is not None:
                  lowerCAmelCase  						=       image

            if mask_image is not None:
                  lowerCAmelCase  						=       mask_image

            if original_image is not None:
                  lowerCAmelCase  						=       original_image

            lowerCAmelCase  						=       pipe_loaded(**lowercase )[0]

            lowerCAmelCase  						=       np.abs(to_np(lowercase ) - to_np(lowercase ) ).max()
            self.assertLess(lowercase      ,						1e-4 )





      def     _snake_case			(       self )		->     Optional[int]:
            lowerCAmelCase  						=       self.get_dummy_components()
            lowerCAmelCase  						=       self.pipeline_class(**lowercase )
            pipe.to(lowercase )
            pipe.set_progress_bar_config(disable=lowercase )

            lowerCAmelCase  						=       self.get_dummy_inputs(lowercase )
            lowerCAmelCase  						=       pipe(**lowercase )[0]

            with tempfile.TemporaryDirectory() as tmpdir:
                  pipe.save_pretrained(lowercase )
                  lowerCAmelCase  						=       self.pipeline_class.from_pretrained(lowercase )
                  pipe_loaded.to(lowercase )
                  pipe_loaded.set_progress_bar_config(disable=lowercase )

            pipe_loaded.unet.set_attn_processor(AttnAddedKVProcessor() )  # For reproducibility tests

            lowerCAmelCase  						=       self.get_dummy_inputs(lowercase )
            lowerCAmelCase  						=       pipe_loaded(**lowercase )[0]

            lowerCAmelCase  						=       np.abs(to_np(lowercase ) - to_np(lowercase ) ).max()
            self.assertLess(lowercase      ,						1e-4 )






