
"""simple docstring"""






import gc
import unittest

import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTextModelWithProjection, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    DDPMScheduler,
    PriorTransformer,
    StableUnCLIPPipeline,
    UNetaDConditionModel,
)
from diffusers.pipelines.stable_diffusion.stable_unclip_image_normalizer import StableUnCLIPImageNormalizer
from diffusers.utils.testing_utils import enable_full_determinism, load_numpy, require_torch_gpu, slow, torch_device

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import (
    PipelineKarrasSchedulerTesterMixin,
    PipelineLatentTesterMixin,
    PipelineTesterMixin,
    assert_mean_pixel_difference,
)


enable_full_determinism()





class 				a__					(						A__   ,							A__   ,							A__   ,							unittest.TestCase       ):
     UpperCAmelCase__         =      StableUnCLIPPipeline
     UpperCAmelCase__         =      TEXT_TO_IMAGE_PARAMS
     UpperCAmelCase__         =      TEXT_TO_IMAGE_BATCH_PARAMS
     UpperCAmelCase__         =      TEXT_TO_IMAGE_IMAGE_PARAMS
     UpperCAmelCase__         =      TEXT_TO_IMAGE_IMAGE_PARAMS

     # TODO(will) Expected attn_bias.stride(1) == 0 to be true, but got false
     UpperCAmelCase__         =      False



     def  lowerCamelCase_						(							self :Tuple   ):







      '''simple docstring'''



      UpperCamelCase_     :					List[str]   	=32
      UpperCamelCase_     :					str   	=embedder_hidden_size

      # prior components

      torch.manual_seed(0   )
      UpperCamelCase_     :					List[str]   	=CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'   )

      torch.manual_seed(0   )
      UpperCamelCase_     :					Any   	=CLIPTextModelWithProjection(
          CLIPTextConfig(
              bos_token_id=0      ,  eos_token_id=2      ,  hidden_size=_lowerCamelCase      ,  projection_dim=_lowerCamelCase      ,  intermediate_size=37      ,  layer_norm_eps=1E-0_5      ,  num_attention_heads=4      ,  num_hidden_layers=5      ,  pad_token_id=1      ,  vocab_size=1_000      ,  )   )

      torch.manual_seed(0   )
      UpperCamelCase_     :					str   	=PriorTransformer(
          num_attention_heads=2      ,  attention_head_dim=12      ,  embedding_dim=_lowerCamelCase      ,  num_layers=1      ,  )

      torch.manual_seed(0   )
      UpperCamelCase_     :					List[str]   	=DDPMScheduler(
          variance_type='fixed_small_log'      ,  prediction_type='sample'      ,  num_train_timesteps=1_000      ,  clip_sample=_lowerCamelCase      ,  clip_sample_range=5.0      ,  beta_schedule='squaredcos_cap_v2'      ,  )

      # regular denoising components

      torch.manual_seed(0   )
      UpperCamelCase_     :					List[str]   	=StableUnCLIPImageNormalizer(embedding_dim=_lowerCamelCase   )
      UpperCamelCase_     :					List[Any]   	=DDPMScheduler(beta_schedule='squaredcos_cap_v2'   )

      torch.manual_seed(0   )
      UpperCamelCase_     :					Optional[int]   	=CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'   )

      torch.manual_seed(0   )
      UpperCamelCase_     :					Dict   	=CLIPTextModel(
          CLIPTextConfig(
              bos_token_id=0      ,  eos_token_id=2      ,  hidden_size=_lowerCamelCase      ,  projection_dim=32      ,  intermediate_size=37      ,  layer_norm_eps=1E-0_5      ,  num_attention_heads=4      ,  num_hidden_layers=5      ,  pad_token_id=1      ,  vocab_size=1_000      ,  )   )

      torch.manual_seed(0   )
      UpperCamelCase_     :					str   	=UNetaDConditionModel(
          sample_size=32      ,  in_channels=4      ,  out_channels=4      ,  down_block_types=('CrossAttnDownBlock2D', 'DownBlock2D')      ,  up_block_types=('UpBlock2D', 'CrossAttnUpBlock2D')      ,  block_out_channels=(32, 64)      ,  attention_head_dim=(2, 4)      ,  class_embed_type='projection'      ,  projection_class_embeddings_input_dim=embedder_projection_dim * 2      ,  cross_attention_dim=_lowerCamelCase      ,  layers_per_block=1      ,  upcast_attention=_lowerCamelCase      ,  use_linear_projection=_lowerCamelCase      ,  )

      torch.manual_seed(0   )
      UpperCamelCase_     :					List[str]   	=DDIMScheduler(
          beta_schedule='scaled_linear'      ,  beta_start=0.00085      ,  beta_end=0.012      ,  prediction_type='v_prediction'      ,  set_alpha_to_one=_lowerCamelCase      ,  steps_offset=1      ,  )

      torch.manual_seed(0   )
      UpperCamelCase_     :					str   	=AutoencoderKL()

      UpperCamelCase_     :					List[str]   	={
          # prior components
          'prior_tokenizer': prior_tokenizer,
          'prior_text_encoder': prior_text_encoder,
          'prior': prior,
          'prior_scheduler': prior_scheduler,
          # image noising components
          'image_normalizer': image_normalizer,
          'image_noising_scheduler': image_noising_scheduler,
          # regular denoising components
          'tokenizer': tokenizer,
          'text_encoder': text_encoder,
          'unet': unet,
          'scheduler': scheduler,
          'vae': vae,
      }

      return components



     def  lowerCamelCase_						(							self :List[str]      ,  _lowerCamelCase :Tuple      ,  _lowerCamelCase :Dict=0   ):







      '''simple docstring'''



      if str(_lowerCamelCase   ).startswith('mps'   ):
       UpperCamelCase_     :					Tuple   	=torch.manual_seed(_lowerCamelCase   )
      else:
       UpperCamelCase_     :					Any   	=torch.Generator(device=_lowerCamelCase   ).manual_seed(_lowerCamelCase   )
      UpperCamelCase_     :					List[str]   	={
          'prompt': 'A painting of a squirrel eating a burger',
          'generator': generator,
          'num_inference_steps': 2,
          'prior_num_inference_steps': 2,
          'output_type': 'numpy',
      }
      return inputs



     def  lowerCamelCase_						(							self :Any   ):







      '''simple docstring'''



      UpperCamelCase_     :					List[str]   	=torch_device == 'cpu'

      self._test_attention_slicing_forward_pass(test_max_difference=_lowerCamelCase   )




     def  lowerCamelCase_						(							self :str   ):







      '''simple docstring'''



      UpperCamelCase_     :					Tuple   	=torch_device in ['cpu', 'mps']

      self._test_inference_batch_single_identical(test_max_difference=_lowerCamelCase   )





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



     def  lowerCamelCase_						(							self :List[str]   ):







      '''simple docstring'''



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



     def  lowerCamelCase_						(							self :int   ):







      '''simple docstring'''



      UpperCamelCase_     :					Tuple   	=load_numpy(
          'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/stable_unclip/stable_unclip_2_1_l_anime_turtle_fp16.npy'   )

      UpperCamelCase_     :					Any   	=StableUnCLIPPipeline.from_pretrained('fusing/stable-unclip-2-1-l'      ,  torch_dtype=torch.floataa   )
      pipe.to(_lowerCamelCase   )
      pipe.set_progress_bar_config(disable=_lowerCamelCase   )
      # stable unclip will oom when integration tests are run on a V100,
      # so turn on memory savings
      pipe.enable_attention_slicing()
      pipe.enable_sequential_cpu_offload()

      UpperCamelCase_     :					Optional[int]   	=torch.Generator(device='cpu'   ).manual_seed(0   )
      UpperCamelCase_     :					Dict   	=pipe('anime turle'      ,  generator=_lowerCamelCase      ,  output_type='np'   )

      UpperCamelCase_     :					Any   	=output.images[0]

      assert image.shape == (768, 768, 3)

      assert_mean_pixel_difference(_lowerCamelCase      ,  _lowerCamelCase   )




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







      '''simple docstring'''



      torch.cuda.empty_cache()
      torch.cuda.reset_max_memory_allocated()
      torch.cuda.reset_peak_memory_stats()

      UpperCamelCase_     :					Any   	=StableUnCLIPPipeline.from_pretrained('fusing/stable-unclip-2-1-l'      ,  torch_dtype=torch.floataa   )
      UpperCamelCase_     :					List[Any]   	=pipe.to(_lowerCamelCase   )
      pipe.set_progress_bar_config(disable=_lowerCamelCase   )
      pipe.enable_attention_slicing()
      pipe.enable_sequential_cpu_offload()

      UpperCamelCase_     :					str   	=pipe(
          'anime turtle'      ,  prior_num_inference_steps=2      ,  num_inference_steps=2      ,  output_type='np'      ,  )

      UpperCamelCase_     :					Union[str, Any]   	=torch.cuda.max_memory_allocated()
      # make sure that less than 7 GB is allocated
      assert mem_bytes < 7 * 10**9


"""simple docstring"""






import copy
import os
from typing import Union

from ...configuration_utils import PretrainedConfig
from ...models.auto.modeling_auto import MODEL_FOR_CAUSAL_LM_MAPPING_NAMES
from ...utils import logging
from ..auto import CONFIG_MAPPING


__SCREAMING_SNAKE_CASE							=     logging.get_logger(__name__)

__SCREAMING_SNAKE_CASE							=     {
    'Salesforce/instruct-blip-flan-t5': 'https://huggingface.co/Salesforce/instruct-blip-flan-t5/resolve/main/config.json',
}





class 				a__					(						A__       ):
     UpperCAmelCase__         =      '''instructblip_vision_model'''



     def __init__(							self :int      ,  _lowerCamelCase :Optional[int]=1_408      ,  _lowerCamelCase :int=6_144      ,  _lowerCamelCase :str=39      ,  _lowerCamelCase :Tuple=16      ,  _lowerCamelCase :List[Any]=224      ,  _lowerCamelCase :int=14      ,  _lowerCamelCase :int="gelu"      ,  _lowerCamelCase :Optional[Any]=1E-6      ,  _lowerCamelCase :str=0.0      ,  _lowerCamelCase :Optional[Any]=1E-1_0      ,  _lowerCamelCase :str=True      ,  **_lowerCamelCase :Optional[Any]      ,  ):







      '''simple docstring'''



      super().__init__(**_lowerCamelCase   )

      UpperCamelCase_     :					Dict   	=hidden_size
      UpperCamelCase_     :					Tuple   	=intermediate_size
      UpperCamelCase_     :					Optional[int]   	=num_hidden_layers
      UpperCamelCase_     :					Union[str, Any]   	=num_attention_heads
      UpperCamelCase_     :					Optional[Any]   	=patch_size
      UpperCamelCase_     :					List[Any]   	=image_size
      UpperCamelCase_     :					Dict   	=initializer_range
      UpperCamelCase_     :					List[Any]   	=attention_dropout
      UpperCamelCase_     :					List[str]   	=layer_norm_eps
      UpperCamelCase_     :					List[str]   	=hidden_act
      UpperCamelCase_     :					Tuple   	=qkv_bias




     @classmethod
     def  lowerCamelCase_						(							cls :Optional[Any]      ,  _lowerCamelCase :Union[str, os.PathLike]      ,  **_lowerCamelCase :Any   ):







      '''simple docstring'''



      cls._set_token_in_kwargs(_lowerCamelCase   )

      UpperCamelCase_	,     UpperCamelCase_     :					List[str]   	=cls.get_config_dict(_lowerCamelCase      ,  **_lowerCamelCase   )

      # get the vision config dict if we are loading from InstructBlipConfig
      if config_dict.get('model_type'   ) == "instructblip":
       UpperCamelCase_     :					Tuple   	=config_dict['vision_config']

      if "model_type" in config_dict and hasattr(cls      ,  'model_type'   ) and config_dict["model_type"] != cls.model_type:
       logger.warning(
           f'''You are using a model of type {config_dict['model_type']} to instantiate a model of type '''
           f'''{cls.model_type}. This is not supported for all configurations of models and can yield errors.'''   )

      return cls.from_dict(_lowerCamelCase      ,  **_lowerCamelCase   )





class 				a__					(						A__       ):
     UpperCAmelCase__         =      '''instructblip_qformer'''



     def __init__(							self :Dict      ,  _lowerCamelCase :List[Any]=30_522      ,  _lowerCamelCase :Optional[int]=768      ,  _lowerCamelCase :Tuple=12      ,  _lowerCamelCase :List[str]=12      ,  _lowerCamelCase :int=3_072      ,  _lowerCamelCase :Tuple="gelu"      ,  _lowerCamelCase :Dict=0.1      ,  _lowerCamelCase :List[Any]=0.1      ,  _lowerCamelCase :Optional[int]=512      ,  _lowerCamelCase :List[Any]=0.02      ,  _lowerCamelCase :int=1E-1_2      ,  _lowerCamelCase :str=0      ,  _lowerCamelCase :Dict="absolute"      ,  _lowerCamelCase :Optional[int]=2      ,  _lowerCamelCase :Union[str, Any]=1_408      ,  **_lowerCamelCase :Tuple      ,  ):







      '''simple docstring'''



      super().__init__(pad_token_id=_lowerCamelCase      ,  **_lowerCamelCase   )

      UpperCamelCase_     :					Optional[Any]   	=vocab_size
      UpperCamelCase_     :					Tuple   	=hidden_size
      UpperCamelCase_     :					Union[str, Any]   	=num_hidden_layers
      UpperCamelCase_     :					List[str]   	=num_attention_heads
      UpperCamelCase_     :					List[str]   	=hidden_act
      UpperCamelCase_     :					Tuple   	=intermediate_size
      UpperCamelCase_     :					Optional[Any]   	=hidden_dropout_prob
      UpperCamelCase_     :					List[str]   	=attention_probs_dropout_prob
      UpperCamelCase_     :					List[str]   	=max_position_embeddings
      UpperCamelCase_     :					List[str]   	=initializer_range
      UpperCamelCase_     :					str   	=layer_norm_eps
      UpperCamelCase_     :					int   	=position_embedding_type
      UpperCamelCase_     :					Tuple   	=cross_attention_frequency
      UpperCamelCase_     :					List[Any]   	=encoder_hidden_size




     @classmethod
     def  lowerCamelCase_						(							cls :Optional[Any]      ,  _lowerCamelCase :Union[str, os.PathLike]      ,  **_lowerCamelCase :List[Any]   ):







      '''simple docstring'''



      cls._set_token_in_kwargs(_lowerCamelCase   )

      UpperCamelCase_	,     UpperCamelCase_     :					Optional[Any]   	=cls.get_config_dict(_lowerCamelCase      ,  **_lowerCamelCase   )

      # get the qformer config dict if we are loading from InstructBlipConfig
      if config_dict.get('model_type'   ) == "instructblip":
       UpperCamelCase_     :					List[Any]   	=config_dict['qformer_config']

      if "model_type" in config_dict and hasattr(cls      ,  'model_type'   ) and config_dict["model_type"] != cls.model_type:
       logger.warning(
           f'''You are using a model of type {config_dict['model_type']} to instantiate a model of type '''
           f'''{cls.model_type}. This is not supported for all configurations of models and can yield errors.'''   )

      return cls.from_dict(_lowerCamelCase      ,  **_lowerCamelCase   )





class 				a__					(						A__       ):
     UpperCAmelCase__         =      '''instructblip'''
     UpperCAmelCase__         =      True



     def __init__(							self :str      ,  _lowerCamelCase :Optional[Any]=None      ,  _lowerCamelCase :List[str]=None      ,  _lowerCamelCase :List[Any]=None      ,  _lowerCamelCase :Union[str, Any]=32      ,  **_lowerCamelCase :Tuple   ):







      '''simple docstring'''



      super().__init__(**_lowerCamelCase   )

      if vision_config is None:
       UpperCamelCase_     :					Optional[int]   	={}
       logger.info('vision_config is None. initializing the InstructBlipVisionConfig with default values.'   )

      if qformer_config is None:
       UpperCamelCase_     :					Union[str, Any]   	={}
       logger.info('qformer_config is None. Initializing the InstructBlipQFormerConfig with default values.'   )

      if text_config is None:
       UpperCamelCase_     :					str   	={}
       logger.info('text_config is None. Initializing the text config with default values (`OPTConfig`).'   )

      UpperCamelCase_     :					Any   	=InstructBlipVisionConfig(**_lowerCamelCase   )
      UpperCamelCase_     :					Optional[Any]   	=InstructBlipQFormerConfig(**_lowerCamelCase   )
      UpperCamelCase_     :					str   	=text_config['model_type'] if 'model_type' in text_config else 'opt'
      UpperCamelCase_     :					int   	=CONFIG_MAPPING[text_model_type](**_lowerCamelCase   )

      UpperCamelCase_     :					Tuple   	=self.text_config.tie_word_embeddings
      UpperCamelCase_     :					Any   	=self.text_config.is_encoder_decoder

      UpperCamelCase_     :					List[Any]   	=num_query_tokens
      UpperCamelCase_     :					Tuple   	=self.vision_config.hidden_size
      UpperCamelCase_     :					Tuple   	=self.text_config.model_type in MODEL_FOR_CAUSAL_LM_MAPPING_NAMES
      UpperCamelCase_     :					int   	=1.0
      UpperCamelCase_     :					List[str]   	=0.02



     @classmethod
     def  lowerCamelCase_						(							cls :Optional[Any]      ,  _lowerCamelCase :InstructBlipVisionConfig      ,  _lowerCamelCase :InstructBlipQFormerConfig      ,  _lowerCamelCase :PretrainedConfig      ,  **_lowerCamelCase :Dict      ,  ):







      '''simple docstring'''



      return cls(
          vision_config=vision_config.to_dict()      ,  qformer_config=qformer_config.to_dict()      ,  text_config=text_config.to_dict()      ,  **_lowerCamelCase      ,  )




     def  lowerCamelCase_						(							self :Optional[Any]   ):







      '''simple docstring'''



      UpperCamelCase_     :					str   	=copy.deepcopy(self.__dict__   )
      UpperCamelCase_     :					Dict   	=self.vision_config.to_dict()
      UpperCamelCase_     :					Tuple   	=self.qformer_config.to_dict()
      UpperCamelCase_     :					List[str]   	=self.text_config.to_dict()
      UpperCamelCase_     :					List[str]   	=self.__class__.model_type
      return output
