






"""simple docstring"""





import gc
import random
import unittest

import numpy as np
import torch

from diffusers import DDIMScheduler, KandinskyVaaPipeline, KandinskyVaaPriorPipeline, UNetaDConditionModel, VQModel
from diffusers.utils import floats_tensor, 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      _lowercase (     UpperCAmelCase__		,  unittest.TestCase						):
      _a   :     List[str]								=   KandinskyVaaPipeline
      _a   :     Optional[Any]								=   [
          "image_embeds",
          "negative_image_embeds",
      ]
      _a   :     List[str]								=   ["image_embeds", "negative_image_embeds"]
      _a   :     Dict								=   [
          "generator",
          "height",
          "width",
          "latents",
          "guidance_scale",
          "num_inference_steps",
          "return_dict",
          "guidance_scale",
          "num_images_per_prompt",
          "output_type",
          "return_dict",
      ]
      _a   :     Any								=   False


      @property
      def _UpperCamelCase						(    self				:	str						):





            """simple docstring"""







            return 3_2


      @property
      def _UpperCamelCase						(    self				:	List[Any]						):





            """simple docstring"""







            return 3_2


      @property
      def _UpperCamelCase						(    self				:	List[Any]						):





            """simple docstring"""







            return self.time_input_dim


      @property
      def _UpperCamelCase						(    self				:	List[Any]						):





            """simple docstring"""







            return self.time_input_dim * 4


      @property
      def _UpperCamelCase						(    self				:	int						):





            """simple docstring"""







            return 1_0_0


      @property
      def _UpperCamelCase						(    self				:	Dict						):





            """simple docstring"""







            torch.manual_seed(0						)

            __snake_case      :    Any  			={
                "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      :    int  			=UNetaDConditionModel(**lowerCamelCase__						)
            return model


      @property
      def _UpperCamelCase						(    self				:	Optional[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 _UpperCamelCase						(    self				:	Tuple						):





            """simple docstring"""







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


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





            """simple docstring"""







            __snake_case      :    Optional[int]  			=self.dummy_unet
            __snake_case      :    List[Any]  			=self.dummy_movq

            __snake_case      :    int  			=DDIMScheduler(
                num_train_timesteps=1_0_0_0							,   beta_schedule='''linear'''							,   beta_start=0.0_0_0_8_5							,   beta_end=0.0_1_2							,   clip_sample=lowerCamelCase__							,   set_alpha_to_one=lowerCamelCase__							,   steps_offset=1							,   prediction_type='''epsilon'''							,   thresholding=lowerCamelCase__							,   )

            __snake_case      :    Union[str, Any]  			={
                "unet": unet,
                "scheduler": scheduler,
                "movq": movq,
            }
            return components


      def _UpperCamelCase						(    self				:	List[str]							,   a				:	Dict							,   a				:	List[Any]=0						):





            """simple docstring"""







            __snake_case      :    Optional[int]  			=floats_tensor((1, self.text_embedder_hidden_size)							,   rng=random.Random(lowerCamelCase__						)						).to(lowerCamelCase__						)
            __snake_case      :    Any  			=floats_tensor((1, self.text_embedder_hidden_size)							,   rng=random.Random(seed + 1						)						).to(
                lowerCamelCase__						)
            if str(lowerCamelCase__						).startswith('''mps'''						):
                  __snake_case      :    Tuple  			=torch.manual_seed(lowerCamelCase__						)
            else:
                  __snake_case      :    Any  			=torch.Generator(device=lowerCamelCase__						).manual_seed(lowerCamelCase__						)
            __snake_case      :    Dict  			={
                "image_embeds": image_embeds,
                "negative_image_embeds": negative_image_embeds,
                "generator": generator,
                "height": 6_4,
                "width": 6_4,
                "guidance_scale": 4.0,
                "num_inference_steps": 2,
                "output_type": "np",
            }
            return inputs







      def _UpperCamelCase						(    self				:	Any						):





            """simple docstring"""







            __snake_case      :    Optional[Any]  			="cpu"

            __snake_case      :    Tuple  			=self.get_dummy_components()

            __snake_case      :    List[Any]  			=self.pipeline_class(**lowerCamelCase__						)
            __snake_case      :    List[Any]  			=pipe.to(lowerCamelCase__						)

            pipe.set_progress_bar_config(disable=lowerCamelCase__						)

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

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

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

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

            __snake_case      :    Tuple  			=np.array(
                [0.6_2_3_7_9_7_6, 1.0, 0.3_6_4_4_1_3_3_2, 1.0, 0.7_0_6_3_9_6_3_4, 0.2_9_8_7_7_1_8_6, 0.8_5_6_5_2_1_2_5, 0.5_2_1_6_8_4_3, 0.5_4_4_5_4_0_4_6]						)

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


      def _UpperCamelCase						(    self				:	Tuple						):





            """simple docstring"""







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







      def _UpperCamelCase						(    self				:	str						):





            """simple docstring"""







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

            __snake_case      :    List[Any]  			=KandinskyVaaPriorPipeline.from_pretrained(
                '''kandinsky-community/kandinsky-2-2-prior'''							,   torch_dtype=torch.floataa						)
            pipe_prior.to(lowerCamelCase__						)

            __snake_case      :    List[Any]  			=KandinskyVaaPipeline.from_pretrained(
                '''kandinsky-community/kandinsky-2-2-decoder'''							,   torch_dtype=torch.floataa						)
            __snake_case      :    Optional[int]  			=pipeline.to(lowerCamelCase__						)
            pipeline.set_progress_bar_config(disable=lowerCamelCase__						)

            __snake_case      :    Tuple  			="red cat, 4k photo"

            __snake_case      :    List[str]  			=torch.Generator(device='''cuda'''						).manual_seed(0						)
            __snake_case      :    List[str]  			=pipe_prior(
                lowerCamelCase__							,   generator=lowerCamelCase__							,   num_inference_steps=5							,   negative_prompt=''''''							,   ).to_tuple()

            __snake_case      :    str  			=torch.Generator(device='''cuda'''						).manual_seed(0						)
            __snake_case      :    str  			=pipeline(
                image_embeds=lowerCamelCase__							,   negative_image_embeds=lowerCamelCase__							,   generator=lowerCamelCase__							,   num_inference_steps=1_0_0							,   output_type='''np'''							,   )

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

            assert image.shape == (5_1_2, 5_1_2, 3)

            assert_mean_pixel_difference(lowerCamelCase__							,   lowerCamelCase__						)










"""simple docstring"""





import json
import os
import unittest

from transformers import MgpstrTokenizer
from transformers.models.mgp_str.tokenization_mgp_str import VOCAB_FILES_NAMES
from transformers.testing_utils import require_tokenizers

from ...test_tokenization_common import TokenizerTesterMixin


@require_tokenizers
class      _lowercase (     lowerCAmelCase		,  unittest.TestCase						):
      _a   :     List[str]								=   MgpstrTokenizer
      _a   :     int								=   False
      _a   :     List[str]								=   {}
      _a   :     Optional[Any]								=   False


      def _UpperCamelCase						(    self				:	str						):





            """simple docstring"""







            super().setUp()

            # fmt: off
            __snake_case      :    Union[str, Any]  			=['''[GO]''', '''[s]''', '''0''', '''1''', '''2''', '''3''', '''4''', '''5''', '''6''', '''7''', '''8''', '''9''', '''a''', '''b''', '''c''', '''d''', '''e''', '''f''', '''g''', '''h''', '''i''', '''j''', '''k''', '''l''', '''m''', '''n''', '''o''', '''p''', '''q''', '''r''', '''s''', '''t''', '''u''', '''v''', '''w''', '''x''', '''y''', '''z''']
            # fmt: on
            __snake_case      :    Optional[Any]  			=dict(zip(a							,   range(len(a						)						)						)						)

            __snake_case      :    int  			=os.path.join(self.tmpdirname							,   VOCAB_FILES_NAMES['''vocab_file''']						)
            with open(self.vocab_file							,   '''w'''							,   encoding='''utf-8'''						) as fp:
                  fp.write(json.dumps(a						) + '''\n'''						)


      def _UpperCamelCase						(    self				:	Any							,   **a				:	int						):





            """simple docstring"""







            return MgpstrTokenizer.from_pretrained(self.tmpdirname							,   **a						)


      def _UpperCamelCase						(    self				:	int							,   a				:	Any						):





            """simple docstring"""







            __snake_case      :    Dict  			='''tester'''
            __snake_case      :    str  			='''tester'''
            return input_text, output_text


      @unittest.skip('''MGP-STR always lower cases letters.'''						)
      def _UpperCamelCase						(    self				:	Any						):





            """simple docstring"""







            pass


      def _UpperCamelCase						(    self				:	Any						):





            """simple docstring"""







            __snake_case      :    List[Any]  			=self.get_tokenizers(do_lower_case=a						)
            for tokenizer in tokenizers:
                  with self.subTest(f'''{tokenizer.__class__.__name__}'''						):
                        __snake_case      :    List[Any]  			='''[SPECIAL_TOKEN]'''

                        tokenizer.add_special_tokens({'''cls_token''': special_token}						)
                        __snake_case      :    Tuple  			=tokenizer.encode([special_token]							,   add_special_tokens=a						)
                        self.assertEqual(len(a						)							,   1						)

                        __snake_case      :    int  			=tokenizer.decode(a							,   skip_special_tokens=a						)
                        self.assertTrue(special_token not in decoded						)


      def _UpperCamelCase						(    self				:	Tuple						):





            """simple docstring"""







            __snake_case      :    Union[str, Any]  			=self.get_tokenizers()
            for tokenizer in tokenizers:
                  with self.subTest(f'''{tokenizer.__class__.__name__}'''						):
                        __snake_case			,				__snake_case      :    Tuple  			=self.get_input_output_texts(a						)

                        __snake_case      :    int  			=tokenizer.tokenize(a						)
                        __snake_case      :    str  			=tokenizer.convert_tokens_to_ids(a						)
                        __snake_case      :    Union[str, Any]  			=tokenizer.encode(a							,   add_special_tokens=a						)
                        self.assertListEqual(a							,   a						)

                        __snake_case      :    Tuple  			=tokenizer.convert_ids_to_tokens(a						)
                        self.assertNotEqual(len(a						)							,   0						)
                        __snake_case      :    Optional[Any]  			=tokenizer.decode(a						)
                        self.assertIsInstance(a							,   a						)

                        self.assertEqual(text_a.replace(''' '''							,   ''''''						)							,   a						)


      @unittest.skip('''MGP-STR tokenizer only handles one sequence.'''						)
      def _UpperCamelCase						(    self				:	List[Any]						):





            """simple docstring"""







            pass







      @unittest.skip('''inputs cannot be pretokenized in MgpstrTokenizer'''						)
      def _UpperCamelCase						(    self				:	List[str]						):





            """simple docstring"""







            pass


