

import gc
import random
import unittest

import numpy as np
import torch
from PIL import Image
from transformers import XLMRobertaTokenizerFast

from diffusers import DDIMScheduler, KandinskyImgaImgPipeline, KandinskyPriorPipeline, UNetaDConditionModel, VQModel
from diffusers.pipelines.kandinsky.text_encoder import MCLIPConfig, MultilingualCLIP
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        a__       (   UpperCAmelCase	,			unittest.TestCase       ):







  """simple docstring"""




  UpperCAmelCase__      : Dict			       =KandinskyImgaImgPipeline
  UpperCAmelCase__      : List[str]			       =["""prompt""", """image_embeds""", """negative_image_embeds""", """image"""]
  UpperCAmelCase__      : Optional[int]			       =[
      """prompt""",
      """negative_prompt""",
      """image_embeds""",
      """negative_image_embeds""",
      """image""",
  ]
  UpperCAmelCase__      : int			       =[
      """generator""",
      """height""",
      """width""",
      """strength""",
      """guidance_scale""",
      """negative_prompt""",
      """num_inference_steps""",
      """return_dict""",
      """guidance_scale""",
      """num_images_per_prompt""",
      """output_type""",
      """return_dict""",
  ]
  UpperCAmelCase__      : List[str]			       =False


  @property
  def  _lowercase      (			self    : int		)					->Union[str, Any]:
        """simple docstring"""







        return 3_2


  @property
  def  _lowercase      (			self    : int		)					->Optional[int]:
        """simple docstring"""







        return 3_2


  @property
  def  _lowercase      (			self    : int		)					->int:
        """simple docstring"""







        return self.time_input_dim


  @property
  def  _lowercase      (			self    : Tuple		)					->Any:
        """simple docstring"""







        return self.time_input_dim * 4


  @property
  def  _lowercase      (			self    : List[str]		)					->List[Any]:
        """simple docstring"""







        return 1_0_0


  @property
  def  _lowercase      (			self    : Any		)					->Optional[Any]:
        """simple docstring"""







        SCREAMING_SNAKE_CASE :      Dict       =   XLMRobertaTokenizerFast.from_pretrained("""YiYiXu/tiny-random-mclip-base"""		)
        return tokenizer


  @property
  def  _lowercase      (			self    : Union[str, Any]		)					->str:
        """simple docstring"""







        torch.manual_seed(0		)
        SCREAMING_SNAKE_CASE :      Union[str, Any]       =   MCLIPConfig(
            numDims=self.cross_attention_dim				,				transformerDimensions=self.text_embedder_hidden_size				,				hidden_size=self.text_embedder_hidden_size				,				intermediate_size=3_7				,				num_attention_heads=4				,				num_hidden_layers=5				,				vocab_size=1_0_0_5				,				)

        SCREAMING_SNAKE_CASE :      int       =   MultilingualCLIP(UpperCAmelCase__		)
        SCREAMING_SNAKE_CASE :      List[Any]       =   text_encoder.eval()

        return text_encoder


  @property
  def  _lowercase      (			self    : Union[str, Any]		)					->Dict:
        """simple docstring"""







        torch.manual_seed(0		)

        SCREAMING_SNAKE_CASE :      Union[str, Any]       =   {
            """in_channels""": 4,
            # Out channels is double in channels because predicts mean and variance
            """out_channels""": 8,
            """addition_embed_type""": """text_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""": """text_image_proj""",
            """cross_attention_dim""": self.cross_attention_dim,
            """attention_head_dim""": 4,
            """resnet_time_scale_shift""": """scale_shift""",
            """class_embed_type""": None,
        }

        SCREAMING_SNAKE_CASE :      Union[str, Any]       =   UNetaDConditionModel(**UpperCAmelCase__		)
        return model


  @property
  def  _lowercase      (			self    : Union[str, Any]		)					->Union[str, Any]:
        """simple docstring"""







        return {
            "block_out_channels": [3_2, 6_4],
            "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": 1_2,
            "out_channels": 3,
            "up_block_types": [
                "AttnUpDecoderBlock2D",
                "UpDecoderBlock2D",
            ],
            "vq_embed_dim": 4,
        }


  @property
  def  _lowercase      (			self    : Optional[Any]		)					->Any:
        """simple docstring"""







        torch.manual_seed(0		)
        SCREAMING_SNAKE_CASE :      Tuple       =   VQModel(**self.dummy_movq_kwargs		)
        return model


  def  _lowercase      (			self    : Any		)					->Optional[int]:
        """simple docstring"""







        SCREAMING_SNAKE_CASE :      Optional[int]       =   self.dummy_text_encoder
        SCREAMING_SNAKE_CASE :      Optional[int]       =   self.dummy_tokenizer
        SCREAMING_SNAKE_CASE :      Optional[int]       =   self.dummy_unet
        SCREAMING_SNAKE_CASE :      List[str]       =   self.dummy_movq

        SCREAMING_SNAKE_CASE :      Any       =   {
            """num_train_timesteps""": 1_0_0_0,
            """beta_schedule""": """linear""",
            """beta_start""": 0.0_00_85,
            """beta_end""": 0.0_12,
            """clip_sample""": False,
            """set_alpha_to_one""": False,
            """steps_offset""": 0,
            """prediction_type""": """epsilon""",
            """thresholding""": False,
        }

        SCREAMING_SNAKE_CASE :      Tuple       =   DDIMScheduler(**UpperCAmelCase__		)

        SCREAMING_SNAKE_CASE :      Optional[Any]       =   {
            """text_encoder""": text_encoder,
            """tokenizer""": tokenizer,
            """unet""": unet,
            """scheduler""": scheduler,
            """movq""": movq,
        }

        return components


  def  _lowercase      (			self    : str				,				UpperCAmelCase__    : int				,				UpperCAmelCase__    : str=0		)					->Union[str, Any]:
        """simple docstring"""







        SCREAMING_SNAKE_CASE :      Dict       =   floats_tensor((1, self.cross_attention_dim)				,				rng=random.Random(UpperCAmelCase__		)		).to(UpperCAmelCase__		)
        SCREAMING_SNAKE_CASE :      str       =   floats_tensor((1, self.cross_attention_dim)				,				rng=random.Random(seed + 1		)		).to(UpperCAmelCase__		)
        # create init_image
        SCREAMING_SNAKE_CASE :      Tuple       =   floats_tensor((1, 3, 6_4, 6_4)				,				rng=random.Random(UpperCAmelCase__		)		).to(UpperCAmelCase__		)
        SCREAMING_SNAKE_CASE :      Union[str, Any]       =   image.cpu().permute(0				,				2				,				3				,				1		)[0]
        SCREAMING_SNAKE_CASE :      Any       =   Image.fromarray(np.uinta(UpperCAmelCase__		)		).convert("""RGB"""		).resize((2_5_6, 2_5_6)		)

        if str(UpperCAmelCase__		).startswith("""mps"""		):
              SCREAMING_SNAKE_CASE :      Any       =   torch.manual_seed(UpperCAmelCase__		)
        else:
              SCREAMING_SNAKE_CASE :      Tuple       =   torch.Generator(device=UpperCAmelCase__		).manual_seed(UpperCAmelCase__		)
        SCREAMING_SNAKE_CASE :      Union[str, Any]       =   {
            """prompt""": """horse""",
            """image""": init_image,
            """image_embeds""": image_embeds,
            """negative_image_embeds""": negative_image_embeds,
            """generator""": generator,
            """height""": 6_4,
            """width""": 6_4,
            """num_inference_steps""": 1_0,
            """guidance_scale""": 7.0,
            """strength""": 0.2,
            """output_type""": """np""",
        }
        return inputs




  def  _lowercase      (			self    : Optional[int]		)					->Union[str, Any]:
        """simple docstring"""







        SCREAMING_SNAKE_CASE :      Tuple       =   """cpu"""

        SCREAMING_SNAKE_CASE :      Any       =   self.get_dummy_components()

        SCREAMING_SNAKE_CASE :      Union[str, Any]       =   self.pipeline_class(**UpperCAmelCase__		)
        SCREAMING_SNAKE_CASE :      List[str]       =   pipe.to(UpperCAmelCase__		)

        pipe.set_progress_bar_config(disable=UpperCAmelCase__		)

        SCREAMING_SNAKE_CASE :      Dict       =   pipe(**self.get_dummy_inputs(UpperCAmelCase__		)		)
        SCREAMING_SNAKE_CASE :      Tuple       =   output.images

        SCREAMING_SNAKE_CASE :      List[str]       =   pipe(
            **self.get_dummy_inputs(UpperCAmelCase__		)				,				return_dict=UpperCAmelCase__				,				)[0]

        SCREAMING_SNAKE_CASE :      Dict       =   image[0, -3:, -3:, -1]
        SCREAMING_SNAKE_CASE :      Optional[Any]       =   image_from_tuple[0, -3:, -3:, -1]

        assert image.shape == (1, 6_4, 6_4, 3)

        SCREAMING_SNAKE_CASE :      Optional[int]       =   np.array(
            [0.61_47_49_43, 0.6_07_35_39, 0.43_30_85_44, 0.5_92_82_69, 0.47_49_35_95, 0.46_75_59_73, 0.4_61_38_38, 0.45_36_87_97, 0.50_11_92_33]		)
        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()}"







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







  """simple docstring"""




  def  _lowercase      (			self    : List[str]		)					->List[str]:
        """simple docstring"""







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




  def  _lowercase      (			self    : Optional[Any]		)					->List[str]:
        """simple docstring"""







        SCREAMING_SNAKE_CASE :      Dict       =   load_numpy(
            """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"""
            """/kandinsky/kandinsky_img2img_frog.npy"""		)

        SCREAMING_SNAKE_CASE :      List[Any]       =   load_image(
            """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main""" """/kandinsky/cat.png"""		)
        SCREAMING_SNAKE_CASE :      Optional[Any]       =   """A red cartoon frog, 4k"""

        SCREAMING_SNAKE_CASE :      Optional[int]       =   KandinskyPriorPipeline.from_pretrained(
            """kandinsky-community/kandinsky-2-1-prior"""				,				torch_dtype=torch.floataa		)
        pipe_prior.to(UpperCAmelCase__		)

        SCREAMING_SNAKE_CASE :      Optional[int]       =   KandinskyImgaImgPipeline.from_pretrained(
            """kandinsky-community/kandinsky-2-1"""				,				torch_dtype=torch.floataa		)
        SCREAMING_SNAKE_CASE :      List[str]       =   pipeline.to(UpperCAmelCase__		)

        pipeline.set_progress_bar_config(disable=UpperCAmelCase__		)

        SCREAMING_SNAKE_CASE :      Optional[int]       =   torch.Generator(device="""cpu"""		).manual_seed(0		)
        SCREAMING_SNAKE_CASE     ,	SCREAMING_SNAKE_CASE :      str       =   pipe_prior(
            UpperCAmelCase__				,				generator=UpperCAmelCase__				,				num_inference_steps=5				,				negative_prompt=""""""				,				).to_tuple()

        SCREAMING_SNAKE_CASE :      int       =   pipeline(
            UpperCAmelCase__				,				image=UpperCAmelCase__				,				image_embeds=UpperCAmelCase__				,				negative_image_embeds=UpperCAmelCase__				,				generator=UpperCAmelCase__				,				num_inference_steps=1_0_0				,				height=7_6_8				,				width=7_6_8				,				strength=0.2				,				output_type="""np"""				,				)

        SCREAMING_SNAKE_CASE :      List[str]       =   output.images[0]

        assert image.shape == (7_6_8, 7_6_8, 3)

        assert_mean_pixel_difference(UpperCAmelCase__				,				UpperCAmelCase__		)







from ...configuration_utils import PretrainedConfig
from ...utils import logging


UpperCAmelCase__ :     List[Any]				    =    logging.get_logger(__name__)

UpperCAmelCase__ :     int				    =    {
    """unc-nlp/lxmert-base-uncased""": """https://huggingface.co/unc-nlp/lxmert-base-uncased/resolve/main/config.json""",
}







class        a__       (   UpperCAmelCase       ):







  """simple docstring"""




  UpperCAmelCase__      : List[str]			       ="""lxmert"""
  UpperCAmelCase__      : str			       ={}


  def __init__(			self    : int				,				UpperCAmelCase__    : Dict=3_0_5_2_2				,				UpperCAmelCase__    : int=7_6_8				,				UpperCAmelCase__    : Union[str, Any]=1_2				,				UpperCAmelCase__    : int=9_5_0_0				,				UpperCAmelCase__    : Union[str, Any]=1_6_0_0				,				UpperCAmelCase__    : Optional[Any]=4_0_0				,				UpperCAmelCase__    : List[str]=3_0_7_2				,				UpperCAmelCase__    : List[Any]="gelu"				,				UpperCAmelCase__    : int=0.1				,				UpperCAmelCase__    : Optional[Any]=0.1				,				UpperCAmelCase__    : Union[str, Any]=5_1_2				,				UpperCAmelCase__    : int=2				,				UpperCAmelCase__    : List[str]=0.02				,				UpperCAmelCase__    : str=1e-12				,				UpperCAmelCase__    : int=9				,				UpperCAmelCase__    : Optional[int]=5				,				UpperCAmelCase__    : int=5				,				UpperCAmelCase__    : List[str]=2_0_4_8				,				UpperCAmelCase__    : Tuple=4				,				UpperCAmelCase__    : List[str]=6.67				,				UpperCAmelCase__    : Optional[int]=True				,				UpperCAmelCase__    : Dict=True				,				UpperCAmelCase__    : Optional[Any]=True				,				UpperCAmelCase__    : Optional[int]=True				,				UpperCAmelCase__    : List[Any]=True				,				UpperCAmelCase__    : Dict=True				,				UpperCAmelCase__    : Optional[Any]=True				,				**UpperCAmelCase__    : Tuple				,				)					->Optional[Any]:
        """simple docstring"""







        SCREAMING_SNAKE_CASE :      List[str]       =   vocab_size
        SCREAMING_SNAKE_CASE :      List[Any]       =   hidden_size
        SCREAMING_SNAKE_CASE :      Any       =   num_attention_heads
        SCREAMING_SNAKE_CASE :      str       =   hidden_act
        SCREAMING_SNAKE_CASE :      Union[str, Any]       =   intermediate_size
        SCREAMING_SNAKE_CASE :      List[str]       =   hidden_dropout_prob
        SCREAMING_SNAKE_CASE :      str       =   attention_probs_dropout_prob
        SCREAMING_SNAKE_CASE :      Dict       =   max_position_embeddings
        SCREAMING_SNAKE_CASE :      Any       =   type_vocab_size
        SCREAMING_SNAKE_CASE :      Optional[int]       =   initializer_range
        SCREAMING_SNAKE_CASE :      Optional[int]       =   layer_norm_eps
        SCREAMING_SNAKE_CASE :      Dict       =   num_qa_labels
        SCREAMING_SNAKE_CASE :      str       =   num_object_labels
        SCREAMING_SNAKE_CASE :      str       =   num_attr_labels
        SCREAMING_SNAKE_CASE :      List[str]       =   l_layers
        SCREAMING_SNAKE_CASE :      Optional[int]       =   x_layers
        SCREAMING_SNAKE_CASE :      List[str]       =   r_layers
        SCREAMING_SNAKE_CASE :      Union[str, Any]       =   visual_feat_dim
        SCREAMING_SNAKE_CASE :      Dict       =   visual_pos_dim
        SCREAMING_SNAKE_CASE :      List[Any]       =   visual_loss_normalizer
        SCREAMING_SNAKE_CASE :      Union[str, Any]       =   task_matched
        SCREAMING_SNAKE_CASE :      int       =   task_mask_lm
        SCREAMING_SNAKE_CASE :      int       =   task_obj_predict
        SCREAMING_SNAKE_CASE :      Optional[int]       =   task_qa
        SCREAMING_SNAKE_CASE :      Any       =   visual_obj_loss
        SCREAMING_SNAKE_CASE :      str       =   visual_attr_loss
        SCREAMING_SNAKE_CASE :      Union[str, Any]       =   visual_feat_loss
        SCREAMING_SNAKE_CASE :      List[Any]       =   {"""vision""": r_layers, """cross_encoder""": x_layers, """language""": l_layers}
        super().__init__(**UpperCAmelCase__		)




