





import gc
import random
import unittest

import numpy as np
import torch
from PIL import Image

from diffusers import (
    DDIMScheduler,
    KandinskyVaaInpaintPipeline,
    KandinskyVaaPriorPipeline,
    UNetaDConditionModel,
    VQModel,
)
from diffusers.utils import floats_tensor, load_image, load_numpy, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..test_pipelines_common import PipelineTesterMixin, assert_mean_pixel_difference


enable_full_determinism()
class SCREAMING_SNAKE_CASE_  (    lowerCAmelCase__				, unittest.TestCase    ):
						__magic_name__:	int       					=				KandinskyVaaInpaintPipeline
						__magic_name__:	int       					=				["""image_embeds""", """negative_image_embeds""", """image""", """mask_image"""]
						__magic_name__:	Dict       					=				[
						    """image_embeds""",
						    """negative_image_embeds""",
						    """image""",
						    """mask_image""",
						]
						__magic_name__:	Any       					=				[
						    """generator""",
						    """height""",
						    """width""",
						    """latents""",
						    """guidance_scale""",
						    """num_inference_steps""",
						    """return_dict""",
						    """guidance_scale""",
						    """num_images_per_prompt""",
						    """output_type""",
						    """return_dict""",
						]
						__magic_name__:	Tuple       					=				False






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







							"""simple docstring"""


							return 32






						@property
						def UpperCAmelCase_ (					self					:					Tuple		)							->			Dict:







							"""simple docstring"""


							return 32






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







							"""simple docstring"""


							return self.time_input_dim






						@property
						def UpperCAmelCase_ (					self					:					Union[str, Any]		)							->			int:







							"""simple docstring"""


							return self.time_input_dim * 4






						@property
						def UpperCAmelCase_ (					self					:					Union[str, Any]		)							->			int:







							"""simple docstring"""


							return 100






						@property
						def UpperCAmelCase_ (					self					:					List[str]		)							->			Tuple:







							"""simple docstring"""


							torch.manual_seed(0		)

							snake_case_						:	str      =		{
							    "in_channels": 9,
							    # Out channels is double in channels because predicts mean and variance
							    "out_channels": 8,
							    "addition_embed_type": "image",
							    "down_block_types": ("ResnetDownsampleBlock2D", "SimpleCrossAttnDownBlock2D"),
							    "up_block_types": ("SimpleCrossAttnUpBlock2D", "ResnetUpsampleBlock2D"),
							    "mid_block_type": "UNetMidBlock2DSimpleCrossAttn",
							    "block_out_channels": (self.block_out_channels_a, self.block_out_channels_a * 2),
							    "layers_per_block": 1,
							    "encoder_hid_dim": self.text_embedder_hidden_size,
							    "encoder_hid_dim_type": "image_proj",
							    "cross_attention_dim": self.cross_attention_dim,
							    "attention_head_dim": 4,
							    "resnet_time_scale_shift": "scale_shift",
							    "class_embed_type": None,
							}

							snake_case_						:	Any      =		UNetaDConditionModel(**UpperCamelCase__		)
							return model






						@property
						def UpperCAmelCase_ (					self					:					Any		)							->			Dict:







							"""simple docstring"""


							return {
							    "block_out_channels": [32, 64],
							    "down_block_types": ["DownEncoderBlock2D", "AttnDownEncoderBlock2D"],
							    "in_channels": 3,
							    "latent_channels": 4,
							    "layers_per_block": 1,
							    "norm_num_groups": 8,
							    "norm_type": "spatial",
							    "num_vq_embeddings": 12,
							    "out_channels": 3,
							    "up_block_types": [
							        "AttnUpDecoderBlock2D",
							        "UpDecoderBlock2D",
							    ],
							    "vq_embed_dim": 4,
							}






						@property
						def UpperCAmelCase_ (					self					:					Union[str, Any]		)							->			Dict:







							"""simple docstring"""


							torch.manual_seed(0		)
							snake_case_						:	int      =		VQModel(**self.dummy_movq_kwargs		)
							return model






						def UpperCAmelCase_ (					self					:					int		)							->			Optional[int]:







							"""simple docstring"""


							snake_case_						:	List[str]      =		self.dummy_unet
							snake_case_						:	int      =		self.dummy_movq

							snake_case_						:	Optional[int]      =		DDIMScheduler(
							    num_train_timesteps=1000						,      beta_schedule='linear'						,      beta_start=0.0_0_0_8_5						,      beta_end=0.0_1_2						,      clip_sample=UpperCamelCase__						,      set_alpha_to_one=UpperCamelCase__						,      steps_offset=1						,      prediction_type='epsilon'						,      thresholding=UpperCamelCase__						,      )

							snake_case_						:	Optional[Any]      =		{
							    "unet": unet,
							    "scheduler": scheduler,
							    "movq": movq,
							}

							return components






						def UpperCAmelCase_ (					self					:					Union[str, Any]						,      _A					:					List[str]						,      _A					:					List[Any]=0		)							->			Any:







							"""simple docstring"""


							snake_case_						:	Optional[int]      =		floats_tensor((1, self.text_embedder_hidden_size)						,      rng=random.Random(UpperCamelCase__		)		).to(UpperCamelCase__		)
							snake_case_						:	int      =		floats_tensor((1, self.text_embedder_hidden_size)						,      rng=random.Random(seed + 1		)		).to(
							    UpperCamelCase__		)
							# create init_image
							snake_case_						:	str      =		floats_tensor((1, 3, 64, 64)						,      rng=random.Random(UpperCamelCase__		)		).to(UpperCamelCase__		)
							snake_case_						:	Union[str, Any]      =		image.cpu().permute(0						,      2						,      3						,      1		)[0]
							snake_case_						:	List[str]      =		Image.fromarray(np.uinta(UpperCamelCase__		)		).convert('RGB'		).resize((256, 256)		)
							# create mask
							snake_case_						:	Dict      =		np.ones((64, 64)						,      dtype=np.floataa		)
							snake_case_						:	Dict      =		0

							if str(UpperCamelCase__		).startswith('mps'		):
								snake_case_						:	Optional[Any]      =		torch.manual_seed(UpperCamelCase__		)
							else:
								snake_case_						:	str      =		torch.Generator(device=UpperCamelCase__		).manual_seed(UpperCamelCase__		)
							snake_case_						:	Optional[Any]      =		{
							    "image": init_image,
							    "mask_image": mask,
							    "image_embeds": image_embeds,
							    "negative_image_embeds": negative_image_embeds,
							    "generator": generator,
							    "height": 64,
							    "width": 64,
							    "num_inference_steps": 2,
							    "guidance_scale": 4.0,
							    "output_type": "np",
							}
							return inputs






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







							"""simple docstring"""


							snake_case_						:	Optional[Any]      =		"cpu"

							snake_case_						:	Optional[Any]      =		self.get_dummy_components()

							snake_case_						:	Optional[Any]      =		self.pipeline_class(**UpperCamelCase__		)
							snake_case_						:	str      =		pipe.to(UpperCamelCase__		)

							pipe.set_progress_bar_config(disable=UpperCamelCase__		)

							snake_case_						:	Any      =		pipe(**self.get_dummy_inputs(UpperCamelCase__		)		)
							snake_case_						:	Dict      =		output.images

							snake_case_						:	Union[str, Any]      =		pipe(
							    **self.get_dummy_inputs(UpperCamelCase__		)						,      return_dict=UpperCamelCase__						,      )[0]

							snake_case_						:	int      =		image[0, -3:, -3:, -1]
							snake_case_						:	Any      =		image_from_tuple[0, -3:, -3:, -1]

							print(F"""image.shape {image.shape}"""		)

							assert image.shape == (1, 64, 64, 3)

							snake_case_						:	Optional[int]      =		np.array(
							    [0.5_0_7_7_5_9_0_3, 0.4_9_5_2_7_1_9_5, 0.4_8_8_2_4_5_4_3, 0.5_0_1_9_2_2_3_7, 0.4_8_6_4_4_9_0_6, 0.4_9_3_7_3_8_1_4, 0.4_7_8_0_5_9_8, 0.4_7_2_3_4_8_2_7, 0.4_8_3_2_7_8_4_8]		)

							assert (
							    np.abs(image_slice.flatten() - expected_slice		).max() < 1E-2
							), F""" expected_slice {expected_slice}, but got {image_slice.flatten()}"""
							assert (
							    np.abs(image_from_tuple_slice.flatten() - expected_slice		).max() < 1E-2
							), F""" expected_slice {expected_slice}, but got {image_from_tuple_slice.flatten()}"""






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







							"""simple docstring"""


							super().test_inference_batch_single_identical(expected_max_diff=3E-3		)
@slow
@require_torch_gpu
class SCREAMING_SNAKE_CASE_  (    unittest.TestCase    ):






						def UpperCAmelCase_ (					self					:					List[Any]		)							->			Optional[int]:







							"""simple docstring"""


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






						def UpperCAmelCase_ (					self					:					List[str]		)							->			Dict:







							"""simple docstring"""


							snake_case_						:	Optional[Any]      =		load_numpy(
							    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
							    '/kandinskyv22/kandinskyv22_inpaint_cat_with_hat_fp16.npy'		)

							snake_case_						:	List[str]      =		load_image(
							    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main' '/kandinsky/cat.png'		)
							snake_case_						:	List[Any]      =		np.ones((768, 768)						,      dtype=np.floataa		)
							snake_case_						:	Any      =		0

							snake_case_						:	Optional[int]      =		"a hat"

							snake_case_						:	Tuple      =		KandinskyVaaPriorPipeline.from_pretrained(
							    'kandinsky-community/kandinsky-2-2-prior'						,      torch_dtype=torch.floataa		)
							pipe_prior.to(UpperCamelCase__		)

							snake_case_						:	int      =		KandinskyVaaInpaintPipeline.from_pretrained(
							    'kandinsky-community/kandinsky-2-2-decoder-inpaint'						,      torch_dtype=torch.floataa		)
							snake_case_						:	Any      =		pipeline.to(UpperCamelCase__		)
							pipeline.set_progress_bar_config(disable=UpperCamelCase__		)

							snake_case_						:	str      =		torch.Generator(device='cpu'		).manual_seed(0		)
							snake_case_						:	Tuple      =		pipe_prior(
							    UpperCamelCase__						,      generator=UpperCamelCase__						,      num_inference_steps=5						,      negative_prompt=''						,      ).to_tuple()

							snake_case_						:	Any      =		pipeline(
							    image=UpperCamelCase__						,      mask_image=UpperCamelCase__						,      image_embeds=UpperCamelCase__						,      negative_image_embeds=UpperCamelCase__						,      generator=UpperCamelCase__						,      num_inference_steps=100						,      height=768						,      width=768						,      output_type='np'						,      )

							snake_case_						:	str      =		output.images[0]

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

							assert_mean_pixel_difference(UpperCamelCase__						,      UpperCamelCase__		)







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__       :      Optional[int]  =   logging.get_logger(__name__)

SCREAMING_SNAKE_CASE__       :      Dict  =   {
    'salesforce/blip2-opt-2.7b': 'https://huggingface.co/salesforce/blip2-opt-2.7b/resolve/main/config.json',
}




class 		UpperCamelCase__	(lowerCAmelCase__     ):

    '''simple docstring'''



    lowerCamelCase_      :	Union[str, Any]	     =      """blip_2_vision_model"""







    def __init__(    self     ,  UpperCamelCase__=1408     ,  UpperCamelCase__=6144     ,  UpperCamelCase__=39     ,  UpperCamelCase__=16     ,  UpperCamelCase__=224     ,  UpperCamelCase__=14     ,  UpperCamelCase__="gelu"     ,  UpperCamelCase__=0.00001     ,  UpperCamelCase__=0.0     ,  UpperCamelCase__=1e-10     ,  UpperCamelCase__=True     ,  **UpperCamelCase__     ,  )     ->      Optional[Any]:
           super().__init__(**UpperCamelCase__     )

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







    @classmethod
    def     _lowercase     (    cls     ,  UpperCamelCase__     ,  **UpperCamelCase__     )     ->      "PretrainedConfig":
           cls._set_token_in_kwargs(UpperCamelCase__     )

           lowerCamelCase ,      lowerCamelCase   :      List[str]			 =					cls.get_config_dict(UpperCamelCase__     ,  **UpperCamelCase__     )

           # get the vision config dict if we are loading from Blip2Config
           if config_dict.get("model_type"     ) == "blip-2":
                  lowerCamelCase   :      Optional[int]			 =					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(UpperCamelCase__     ,  **UpperCamelCase__     )




class 		UpperCamelCase__	(lowerCAmelCase__     ):

    '''simple docstring'''



    lowerCamelCase_      :	Dict	     =      """blip_2_qformer"""







    def __init__(    self     ,  UpperCamelCase__=3_0522     ,  UpperCamelCase__=768     ,  UpperCamelCase__=12     ,  UpperCamelCase__=12     ,  UpperCamelCase__=3072     ,  UpperCamelCase__="gelu"     ,  UpperCamelCase__=0.1     ,  UpperCamelCase__=0.1     ,  UpperCamelCase__=512     ,  UpperCamelCase__=0.02     ,  UpperCamelCase__=1e-12     ,  UpperCamelCase__=0     ,  UpperCamelCase__="absolute"     ,  UpperCamelCase__=2     ,  UpperCamelCase__=1408     ,  **UpperCamelCase__     ,  )     ->      int:
           super().__init__(pad_token_id=UpperCamelCase__     ,  **UpperCamelCase__     )

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







    @classmethod
    def     _lowercase     (    cls     ,  UpperCamelCase__     ,  **UpperCamelCase__     )     ->      "PretrainedConfig":
           cls._set_token_in_kwargs(UpperCamelCase__     )

           lowerCamelCase ,      lowerCamelCase   :      str			 =					cls.get_config_dict(UpperCamelCase__     ,  **UpperCamelCase__     )

           # get the qformer config dict if we are loading from Blip2Config
           if config_dict.get("model_type"     ) == "blip-2":
                  lowerCamelCase   :      int			 =					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(UpperCamelCase__     ,  **UpperCamelCase__     )




class 		UpperCamelCase__	(lowerCAmelCase__     ):

    '''simple docstring'''



    lowerCamelCase_      :	List[str]	     =      """blip-2"""
    lowerCamelCase_      :	int	     =      True







    def __init__(    self     ,  UpperCamelCase__=None     ,  UpperCamelCase__=None     ,  UpperCamelCase__=None     ,  UpperCamelCase__=32     ,  **UpperCamelCase__     )     ->      str:
           super().__init__(**UpperCamelCase__     )

           if vision_config is None:
                  lowerCamelCase   :      List[Any]			 =					{}
                  logger.info("vision_config is None. initializing the Blip2VisionConfig with default values."     )

           if qformer_config is None:
                  lowerCamelCase   :      List[Any]			 =					{}
                  logger.info("qformer_config is None. Initializing the Blip2QFormerConfig with default values."     )

           if text_config is None:
                  lowerCamelCase   :      Any			 =					{}
                  logger.info("text_config is None. Initializing the text config with default values (`OPTConfig`)."     )

           lowerCamelCase   :      Optional[int]			 =					BlipaVisionConfig(**UpperCamelCase__     )
           lowerCamelCase   :      str			 =					BlipaQFormerConfig(**UpperCamelCase__     )
           lowerCamelCase   :      List[str]			 =					text_config["model_type"] if "model_type" in text_config else "opt"
           lowerCamelCase   :      str			 =					CONFIG_MAPPING[text_model_type](**UpperCamelCase__     )

           lowerCamelCase   :      Optional[Any]			 =					self.text_config.tie_word_embeddings
           lowerCamelCase   :      int			 =					self.text_config.is_encoder_decoder

           lowerCamelCase   :      Optional[Any]			 =					num_query_tokens
           lowerCamelCase   :      int			 =					self.vision_config.hidden_size
           lowerCamelCase   :      Tuple			 =					self.text_config.model_type in MODEL_FOR_CAUSAL_LM_MAPPING_NAMES
           lowerCamelCase   :      Dict			 =					1.0
           lowerCamelCase   :      List[Any]			 =					0.02







    @classmethod
    def     _lowercase     (    cls     ,  UpperCamelCase__     ,  UpperCamelCase__     ,  UpperCamelCase__     ,  **UpperCamelCase__     ,  )     ->      str:
           return cls(
               vision_config=vision_config.to_dict()     ,  qformer_config=qformer_config.to_dict()     ,  text_config=text_config.to_dict()     ,  **UpperCamelCase__     ,  )







    def     _lowercase     (    self     )     ->      Optional[Any]:
           lowerCamelCase   :      Tuple			 =					copy.deepcopy(self.__dict__     )
           lowerCamelCase   :      Tuple			 =					self.vision_config.to_dict()
           lowerCamelCase   :      int			 =					self.qformer_config.to_dict()
           lowerCamelCase   :      Optional[Any]			 =					self.text_config.to_dict()
           lowerCamelCase   :      int			 =					self.__class__.model_type
           return output






