






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 			snake_case__   (							lowercase_	,    unittest.TestCase):

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



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






            '''simple docstring'''





            return 32



       @property
       def 					__lowercase						(							self    ) ->      Any:






            '''simple docstring'''





            return 32



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






            '''simple docstring'''





            return self.time_input_dim



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






            '''simple docstring'''





            return self.time_input_dim * 4



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






            '''simple docstring'''





            return 1_00



       @property
       def 					__lowercase						(							self    ) ->      Any:






            '''simple docstring'''





            __snake_case			:List[str]          =						XLMRobertaTokenizerFast.from_pretrained("""YiYiXu/tiny-random-mclip-base"""    )
            return tokenizer



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






            '''simple docstring'''





            torch.manual_seed(0    )
            __snake_case			:Dict          =						MCLIPConfig(
                numDims=self.cross_attention_dim			,						transformerDimensions=self.text_embedder_hidden_size			,						hidden_size=self.text_embedder_hidden_size			,						intermediate_size=37			,						num_attention_heads=4			,						num_hidden_layers=5			,						vocab_size=10_05			,						)

            __snake_case			:Any          =						MultilingualCLIP(UpperCamelCase__    )
            __snake_case			:Optional[int]          =						text_encoder.eval()

            return text_encoder



       @property
       def 					__lowercase						(							self    ) ->      Any:






            '''simple docstring'''





            torch.manual_seed(0    )

            __snake_case			:str          =						{
                """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,
            }

            __snake_case			:int          =						UNetaDConditionModel(**UpperCamelCase__    )
            return model



       @property
       def 					__lowercase						(							self    ) ->      Tuple:






            '''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    ) ->      Tuple:






            '''simple docstring'''





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



       def 					__lowercase						(							self    ) ->      int:






            '''simple docstring'''





            __snake_case			:Optional[int]          =						self.dummy_text_encoder
            __snake_case			:Dict          =						self.dummy_tokenizer
            __snake_case			:Union[str, Any]          =						self.dummy_unet
            __snake_case			:Optional[int]          =						self.dummy_movq

            __snake_case			: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			:List[Any]          =						DDIMScheduler(**UpperCamelCase__    )

            __snake_case			:List[str]          =						{
                """text_encoder""": text_encoder,
                """tokenizer""": tokenizer,
                """unet""": unet,
                """scheduler""": scheduler,
                """movq""": movq,
            }

            return components



       def 					__lowercase						(							self			,						a__			,						a__=0    ) ->      Any:






            '''simple docstring'''





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

            if str(UpperCamelCase__    ).startswith("""mps"""    ):
                 __snake_case			:int          =						torch.manual_seed(UpperCamelCase__    )
            else:
                 __snake_case			:Union[str, Any]          =						torch.Generator(device=UpperCamelCase__    ).manual_seed(UpperCamelCase__    )
            __snake_case			:List[Any]          =						{
                """prompt""": """horse""",
                """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    ) ->      List[Any]:






            '''simple docstring'''





            __snake_case			:int          =						"""cpu"""

            __snake_case			:Tuple          =						self.get_dummy_components()

            __snake_case			:List[str]          =						self.pipeline_class(**UpperCamelCase__    )
            __snake_case			:str          =						pipe.to(UpperCamelCase__    )

            pipe.set_progress_bar_config(disable=UpperCamelCase__    )

            __snake_case			:List[str]          =						pipe(**self.get_dummy_inputs(UpperCamelCase__    )    )
            __snake_case			:Union[str, Any]          =						output.images

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

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

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

            __snake_case			: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 			snake_case__   (							unittest.TestCase):

       '''simple docstring'''



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






            '''simple docstring'''





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



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






            '''simple docstring'''





            __snake_case			:Union[str, Any]          =						load_numpy(
                """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"""
                """/kandinsky/kandinsky_img2img_frog.npy"""    )

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

            __snake_case			:Union[str, Any]          =						KandinskyPriorPipeline.from_pretrained(
                """kandinsky-community/kandinsky-2-1-prior"""			,						torch_dtype=torch.floataa    )
            pipe_prior.to(UpperCamelCase__    )

            __snake_case			:str          =						KandinskyImgaImgPipeline.from_pretrained(
                """kandinsky-community/kandinsky-2-1"""			,						torch_dtype=torch.floataa    )
            __snake_case			:int          =						pipeline.to(UpperCamelCase__    )

            pipeline.set_progress_bar_config(disable=UpperCamelCase__    )

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

            __snake_case			:int          =						pipeline(
                UpperCamelCase__			,						image=UpperCamelCase__			,						image_embeds=UpperCamelCase__			,						negative_image_embeds=UpperCamelCase__			,						generator=UpperCamelCase__			,						num_inference_steps=1_00			,						height=7_68			,						width=7_68			,						strength=0.2			,						output_type="""np"""			,						)

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

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

            assert_mean_pixel_difference(UpperCamelCase__			,						UpperCamelCase__    )










from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_tokenizers_available, is_torch_available


lowerCamelCase__							=     {
    """configuration_canine""": ["""CANINE_PRETRAINED_CONFIG_ARCHIVE_MAP""", """CanineConfig"""],
    """tokenization_canine""": ["""CanineTokenizer"""],
}

try:
     if not is_torch_available():
          raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
     pass
else:
     lowerCamelCase__							=     [
         """CANINE_PRETRAINED_MODEL_ARCHIVE_LIST""",
         """CanineForMultipleChoice""",
         """CanineForQuestionAnswering""",
         """CanineForSequenceClassification""",
         """CanineForTokenClassification""",
         """CanineLayer""",
         """CanineModel""",
         """CaninePreTrainedModel""",
         """load_tf_weights_in_canine""",
     ]


if TYPE_CHECKING:
     from .configuration_canine import CANINE_PRETRAINED_CONFIG_ARCHIVE_MAP, CanineConfig
     from .tokenization_canine import CanineTokenizer

     try:
          if not is_torch_available():
               raise OptionalDependencyNotAvailable()
     except OptionalDependencyNotAvailable:
          pass
     else:
          from .modeling_canine import (
              CANINE_PRETRAINED_MODEL_ARCHIVE_LIST,
              CanineForMultipleChoice,
              CanineForQuestionAnswering,
              CanineForSequenceClassification,
              CanineForTokenClassification,
              CanineLayer,
              CanineModel,
              CaninePreTrainedModel,
              load_tf_weights_in_canine,
          )


else:
     import sys

     lowerCamelCase__							=     _LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)


