






import gc
import random
import unittest

import numpy as np
import torch
from PIL import Image

from diffusers import (
    DDIMScheduler,
    KandinskyVaaImgaImgPipeline,
    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       snake_case__ (				lowercase_						,     unittest.TestCase):

       '''simple docstring'''
       lowerCamelCase				:					List[Any]								=			KandinskyVaaImgaImgPipeline
       lowerCamelCase				:					Optional[Any]								=			["image_embeds", "negative_image_embeds", "image"]
       lowerCamelCase				:					Optional[Any]								=			[
           "image_embeds",
           "negative_image_embeds",
           "image",
       ]
       lowerCamelCase				:					Tuple								=			[
           "generator",
           "height",
           "width",
           "strength",
           "guidance_scale",
           "num_inference_steps",
           "return_dict",
           "guidance_scale",
           "num_images_per_prompt",
           "output_type",
           "return_dict",
       ]
       lowerCamelCase				:					Optional[Any]								=			False



       @property
       def        __lowercase    (				self     )			->   Optional[Any]:






            '''simple docstring'''





            return 32



       @property
       def        __lowercase    (				self     )			->   Any:






            '''simple docstring'''





            return 32



       @property
       def        __lowercase    (				self     )			->   Tuple:






            '''simple docstring'''





            return self.time_input_dim



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






            '''simple docstring'''





            return self.time_input_dim * 4



       @property
       def        __lowercase    (				self     )			->   str:






            '''simple docstring'''





            return 1_00



       @property
       def        __lowercase    (				self     )			->   Optional[int]:






            '''simple docstring'''





            torch.manual_seed(0     )

            __snake_case      :List[str] 							=    {
                """in_channels""": 4,
                # 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      :Optional[int] 							=    UNetaDConditionModel(**a__     )
            return model



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






            '''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        __lowercase    (				self     )			->   List[str]:






            '''simple docstring'''





            torch.manual_seed(0     )
            __snake_case      :List[str] 							=    VQModel(**self.dummy_movq_kwargs     )
            return model



       def        __lowercase    (				self     )			->   Tuple:






            '''simple docstring'''





            __snake_case      :List[str] 							=    self.dummy_unet
            __snake_case      :Any 							=    self.dummy_movq

            __snake_case      :List[str] 							=    {
                """num_train_timesteps""": 10_00,
                """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,
            }

            __snake_case      :str 							=    DDIMScheduler(**a__     )

            __snake_case      :List[Any] 							=    {
                """unet""": unet,
                """scheduler""": scheduler,
                """movq""": movq,
            }

            return components



       def        __lowercase    (				self      ,    a__      ,    a__=0     )			->   List[str]:






            '''simple docstring'''





            __snake_case      :List[Any] 							=    floats_tensor((1, self.text_embedder_hidden_size)      ,    rng=random.Random(a__     )     ).to(a__     )
            __snake_case      :Optional[Any] 							=    floats_tensor((1, self.text_embedder_hidden_size)      ,    rng=random.Random(seed + 1     )     ).to(
                a__     )
            # create init_image
            __snake_case      :Union[str, Any] 							=    floats_tensor((1, 3, 64, 64)      ,    rng=random.Random(a__     )     ).to(a__     )
            __snake_case      :Dict 							=    image.cpu().permute(0      ,    2      ,    3      ,    1     )[0]
            __snake_case      :Union[str, Any] 							=    Image.fromarray(np.uinta(a__     )     ).convert("""RGB"""     ).resize((2_56, 2_56)     )

            if str(a__     ).startswith("""mps"""     ):
                 __snake_case      :Any 							=    torch.manual_seed(a__     )
            else:
                 __snake_case      :Any 							=    torch.Generator(device=a__     ).manual_seed(a__     )
            __snake_case      :Any 							=    {
                """image""": init_image,
                """image_embeds""": image_embeds,
                """negative_image_embeds""": negative_image_embeds,
                """generator""": generator,
                """height""": 64,
                """width""": 64,
                """num_inference_steps""": 10,
                """guidance_scale""": 7.0,
                """strength""": 0.2,
                """output_type""": """np""",
            }
            return inputs



       def        __lowercase    (				self     )			->   Dict:






            '''simple docstring'''





            __snake_case      :List[str] 							=    """cpu"""

            __snake_case      :Optional[int] 							=    self.get_dummy_components()

            __snake_case      :Dict 							=    self.pipeline_class(**a__     )
            __snake_case      :Tuple 							=    pipe.to(a__     )

            pipe.set_progress_bar_config(disable=a__     )

            __snake_case      :Optional[int] 							=    pipe(**self.get_dummy_inputs(a__     )     )
            __snake_case      :Any 							=    output.images

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

            __snake_case      :str 							=    image[0, -3:, -3:, -1]
            __snake_case      :str 							=    image_from_tuple[0, -3:, -3:, -1]

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

            __snake_case      :Optional[int] 							=    np.array(
                [0.6_19_97_78, 0.63_98_44_06, 0.46_14_57_85, 0.62_94_49_84, 0.5_62_22_15, 0.47_30_61_32, 0.47_44_14_56, 0.4_60_76_06, 0.48_71_92_63]     )
            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       snake_case__ (				unittest.TestCase):

       '''simple docstring'''



       def        __lowercase    (				self     )			->   List[str]:






            '''simple docstring'''





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



       def        __lowercase    (				self     )			->   Dict:






            '''simple docstring'''





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

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

            __snake_case      :Optional[Any] 							=    KandinskyVaaPriorPipeline.from_pretrained(
                """kandinsky-community/kandinsky-2-2-prior"""      ,    torch_dtype=torch.floataa     )
            pipe_prior.to(a__     )

            __snake_case      :Tuple 							=    KandinskyVaaImgaImgPipeline.from_pretrained(
                """kandinsky-community/kandinsky-2-2-decoder"""      ,    torch_dtype=torch.floataa     )
            __snake_case      :Optional[Any] 							=    pipeline.to(a__     )

            pipeline.set_progress_bar_config(disable=a__     )

            __snake_case      :Tuple 							=    torch.Generator(device="""cpu"""     ).manual_seed(0     )
            __snake_case							,			__snake_case      :int 							=    pipe_prior(
                a__      ,    generator=a__      ,    num_inference_steps=5      ,    negative_prompt=""""""      ,    ).to_tuple()

            __snake_case      :List[Any] 							=    pipeline(
                image=a__      ,    image_embeds=a__      ,    negative_image_embeds=a__      ,    generator=a__      ,    num_inference_steps=1_00      ,    height=7_68      ,    width=7_68      ,    strength=0.2      ,    output_type="""np"""      ,    )

            __snake_case      :Optional[Any] 							=    output.images[0]

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

            assert_mean_pixel_difference(a__      ,    a__     )










from ...configuration_utils import PretrainedConfig


lowerCamelCase__       		=  {
    """google/tapas-base-finetuned-sqa""": (
        """https://huggingface.co/google/tapas-base-finetuned-sqa/resolve/main/config.json"""
    ),
    """google/tapas-base-finetuned-wtq""": (
        """https://huggingface.co/google/tapas-base-finetuned-wtq/resolve/main/config.json"""
    ),
    """google/tapas-base-finetuned-wikisql-supervised""": (
        """https://huggingface.co/google/tapas-base-finetuned-wikisql-supervised/resolve/main/config.json"""
    ),
    """google/tapas-base-finetuned-tabfact""": (
        """https://huggingface.co/google/tapas-base-finetuned-tabfact/resolve/main/config.json"""
    ),
}




class       snake_case__ (				lowercase_):

       '''simple docstring'''
       lowerCamelCase				:					Tuple								=			"tapas"



       def __init__(				self      ,    a__=3_05_22      ,    a__=7_68      ,    a__=12      ,    a__=12      ,    a__=30_72      ,    a__="gelu"      ,    a__=0.1      ,    a__=0.1      ,    a__=10_24      ,    a__=[3, 2_56, 2_56, 2, 2_56, 2_56, 10]      ,    a__=0.02      ,    a__=1e-12      ,    a__=0      ,    a__=10.0      ,    a__=0      ,    a__=1.0      ,    a__=None      ,    a__=1.0      ,    a__=False      ,    a__=None      ,    a__=1.0      ,    a__=1.0      ,    a__=False      ,    a__=False      ,    a__="ratio"      ,    a__=None      ,    a__=None      ,    a__=64      ,    a__=32      ,    a__=False      ,    a__=True      ,    a__=False      ,    a__=False      ,    a__=True      ,    a__=False      ,    a__=None      ,    a__=None      ,    **a__      ,    )			->   Any:






            '''simple docstring'''





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

            # BERT hyperparameters (with updated max_position_embeddings and type_vocab_sizes)
            __snake_case      :Optional[Any] 							=    vocab_size
            __snake_case      :Tuple 							=    hidden_size
            __snake_case      :Tuple 							=    num_hidden_layers
            __snake_case      :Any 							=    num_attention_heads
            __snake_case      :List[str] 							=    hidden_act
            __snake_case      :List[str] 							=    intermediate_size
            __snake_case      :int 							=    hidden_dropout_prob
            __snake_case      :List[Any] 							=    attention_probs_dropout_prob
            __snake_case      :Any 							=    max_position_embeddings
            __snake_case      :Optional[int] 							=    type_vocab_sizes
            __snake_case      :Tuple 							=    initializer_range
            __snake_case      :Any 							=    layer_norm_eps

            # Fine-tuning task hyperparameters
            __snake_case      :Optional[int] 							=    positive_label_weight
            __snake_case      :Optional[Any] 							=    num_aggregation_labels
            __snake_case      :Dict 							=    aggregation_loss_weight
            __snake_case      :str 							=    use_answer_as_supervision
            __snake_case      :int 							=    answer_loss_importance
            __snake_case      :List[Any] 							=    use_normalized_answer_loss
            __snake_case      :List[str] 							=    huber_loss_delta
            __snake_case      :Any 							=    temperature
            __snake_case      :Optional[int] 							=    aggregation_temperature
            __snake_case      :List[str] 							=    use_gumbel_for_cells
            __snake_case      :Optional[int] 							=    use_gumbel_for_aggregation
            __snake_case      :Optional[Any] 							=    average_approximation_function
            __snake_case      :Union[str, Any] 							=    cell_selection_preference
            __snake_case      :Tuple 							=    answer_loss_cutoff
            __snake_case      :Union[str, Any] 							=    max_num_rows
            __snake_case      :Optional[int] 							=    max_num_columns
            __snake_case      :str 							=    average_logits_per_cell
            __snake_case      :Union[str, Any] 							=    select_one_column
            __snake_case      :Optional[int] 							=    allow_empty_column_selection
            __snake_case      :Optional[Any] 							=    init_cell_selection_weights_to_zero
            __snake_case      :Tuple 							=    reset_position_index_per_cell
            __snake_case      :Any 							=    disable_per_token_loss

            # Aggregation hyperparameters
            __snake_case      :List[str] 							=    aggregation_labels
            __snake_case      :int 							=    no_aggregation_label_index

            if isinstance(self.aggregation_labels      ,    a__     ):
                 __snake_case      :Any 							=    {int(a__     ): v for k, v in aggregation_labels.items()}


