





'''simple docstring'''

import os
import shutil
import tempfile
from unittest import TestCase
from unittest.mock import patch

import numpy as np
from datasets import Dataset

from transformers.models.realm.configuration_realm import RealmConfig
from transformers.models.realm.retrieval_realm import _REALM_BLOCK_RECORDS_FILENAME, RealmRetriever
from transformers.models.realm.tokenization_realm import VOCAB_FILES_NAMES, RealmTokenizer






class    UpperCAmelCase	(       UpperCamelCase__				):
 def      UpperCAmelCase_							(				self      :Union[str, Any]							)->							List[str]:
     A__				=			tempfile.mkdtemp()
     A__				=			5

     # Realm tok
     A__				=			[
         "[UNK]",
         "[CLS]",
         "[SEP]",
         "[PAD]",
         "[MASK]",
         "test",
         "question",
         "this",
         "is",
         "the",
         "first",
         "second",
         "third",
         "fourth",
         "fifth",
         "record",
         "want",
         "##want",
         "##ed",
         "wa",
         "un",
         "runn",
         "##ing",
         ",",
         "low",
         "lowest",
     ]
     A__				=			os.path.join(self.tmpdirname						,       "realm_tokenizer"							)
     os.makedirs(lowercase_						,       exist_ok=lowercase_							)
     A__				=			os.path.join(lowercase_						,       VOCAB_FILES_NAMES["vocab_file"]							)
     with open(self.vocab_file						,       "w"						,       encoding="utf-8"							) as vocab_writer:
         vocab_writer.write("".join([x + "\n" for x in vocab_tokens]							)							)

     A__				=			os.path.join(self.tmpdirname						,       "realm_block_records"							)
     os.makedirs(lowercase_						,       exist_ok=lowercase_							)
 def      UpperCAmelCase_							(				self      :Any							)->							RealmTokenizer:
     return RealmTokenizer.from_pretrained(os.path.join(self.tmpdirname						,       "realm_tokenizer"							)							)
 def      UpperCAmelCase_							(				self      :int							)->							int:
     shutil.rmtree(self.tmpdirname							)
 def      UpperCAmelCase_							(				self      :Tuple							)->							Any:
     A__				=			RealmConfig(num_block_records=self.num_block_records							)
     return config
 def      UpperCAmelCase_							(				self      :Tuple							)->							List[str]:
     A__				=			Dataset.from_dict(
         {
             "id": ["0", "1"],
             "question": ["foo", "bar"],
             "answers": [["Foo", "Bar"], ["Bar"]],
         }							)
     return dataset
 def      UpperCAmelCase_							(				self      :int							)->							str:
     A__				=			np.array(
         [
             B"This is the first record",
             B"This is the second record",
             B"This is the third record",
             B"This is the fourth record",
             B"This is the fifth record",
             B"This is a longer longer longer record",
         ]						,       dtype=lowercase_						,       )
     return block_records
 def      UpperCAmelCase_							(				self      :int							)->							Union[str, Any]:
     A__				=			RealmRetriever(
         block_records=self.get_dummy_block_records()						,       tokenizer=self.get_tokenizer()						,       )
     return retriever
 def      UpperCAmelCase_							(				self      :Optional[Any]							)->							Optional[int]:
     A__				=			self.get_config()
     A__				=			self.get_dummy_retriever()
     A__				=			retriever.tokenizer

     A__				=			np.array([0, 3]						,       dtype="long"							)
     A__				=			tokenizer(["Test question"]							).input_ids
     A__				=			tokenizer(
         ["the fourth"]						,       add_special_tokens=lowercase_						,       return_token_type_ids=lowercase_						,       return_attention_mask=lowercase_						,       ).input_ids
     A__				=			config.reader_seq_len

     A__,	A__,	A__,	A__				=			retriever(
         lowercase_						,       lowercase_						,       answer_ids=lowercase_						,       max_length=lowercase_						,       return_tensors="np"							)

     self.assertEqual(len(lowercase_							)						,       2							)
     self.assertEqual(len(lowercase_							)						,       2							)
     self.assertEqual(len(lowercase_							)						,       2							)
     self.assertEqual(concat_inputs.input_ids.shape						,       (2, 10)							)
     self.assertEqual(concat_inputs.attention_mask.shape						,       (2, 10)							)
     self.assertEqual(concat_inputs.token_type_ids.shape						,       (2, 10)							)
     self.assertEqual(concat_inputs.special_tokens_mask.shape						,       (2, 10)							)
     self.assertEqual(
         tokenizer.convert_ids_to_tokens(concat_inputs.input_ids[0]							)						,       ["[CLS]", "test", "question", "[SEP]", "this", "is", "the", "first", "record", "[SEP]"]						,       )
     self.assertEqual(
         tokenizer.convert_ids_to_tokens(concat_inputs.input_ids[1]							)						,       ["[CLS]", "test", "question", "[SEP]", "this", "is", "the", "fourth", "record", "[SEP]"]						,       )
 def      UpperCAmelCase_							(				self      :List[Any]							)->							Any:
     A__				=			self.get_config()
     A__				=			self.get_dummy_retriever()
     A__				=			retriever.tokenizer

     A__				=			np.array([0, 3, 5]						,       dtype="long"							)
     A__				=			tokenizer(["Test question"]							).input_ids
     A__				=			tokenizer(
         ["the fourth", "longer longer"]						,       add_special_tokens=lowercase_						,       return_token_type_ids=lowercase_						,       return_attention_mask=lowercase_						,       ).input_ids
     A__				=			config.reader_seq_len

     A__,	A__,	A__,	A__				=			retriever(
         lowercase_						,       lowercase_						,       answer_ids=lowercase_						,       max_length=lowercase_						,       return_tensors="np"							)

     self.assertEqual([False, True, True]						,       lowercase_							)
     self.assertEqual([[-1, -1, -1], [6, -1, -1], [6, 7, 8]]						,       lowercase_							)
     self.assertEqual([[-1, -1, -1], [7, -1, -1], [7, 8, 9]]						,       lowercase_							)



 def      UpperCAmelCase_							(				self      :Dict							)->							Union[str, Any]:
     A__				=			self.get_dummy_retriever()
     retriever.save_pretrained(os.path.join(self.tmpdirname						,       "realm_block_records"							)							)

     # Test local path
     A__				=			retriever.from_pretrained(os.path.join(self.tmpdirname						,       "realm_block_records"							)							)
     self.assertEqual(retriever.block_records[0]						,       B"This is the first record"							)

     # Test mocked remote path
     with patch("transformers.models.realm.retrieval_realm.hf_hub_download"							) as mock_hf_hub_download:
         A__				=			os.path.join(
             os.path.join(self.tmpdirname						,       "realm_block_records"							)						,       _REALM_BLOCK_RECORDS_FILENAME							)
         A__				=			RealmRetriever.from_pretrained("google/realm-cc-news-pretrained-openqa"							)

     self.assertEqual(retriever.block_records[0]						,       B"This is the first record"							)










'''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    UpperCAmelCase	(       UpperCamelCase__			,       unittest.TestCase				):
 __lowercase   		=				KandinskyVaaPipeline
 __lowercase   		=				[
     """image_embeds""",
     """negative_image_embeds""",
 ]
 __lowercase   		=				["""image_embeds""", """negative_image_embeds"""]
 __lowercase   		=				[
     """generator""",
     """height""",
     """width""",
     """latents""",
     """guidance_scale""",
     """num_inference_steps""",
     """return_dict""",
     """guidance_scale""",
     """num_images_per_prompt""",
     """output_type""",
     """return_dict""",
 ]
 __lowercase   		=				False
 @property
 def      UpperCAmelCase_							(				self      :Optional[Any]							)->							str:
     return 32
 @property
 def      UpperCAmelCase_							(				self      :int							)->							List[Any]:
     return 32
 @property
 def      UpperCAmelCase_							(				self      :Optional[Any]							)->							Dict:
     return self.time_input_dim
 @property
 def      UpperCAmelCase_							(				self      :Any							)->							Union[str, Any]:
     return self.time_input_dim * 4
 @property
 def      UpperCAmelCase_							(				self      :Any							)->							List[Any]:
     return 1_00
 @property
 def      UpperCAmelCase_							(				self      :Union[str, Any]							)->							Union[str, Any]:
     torch.manual_seed(0							)

     A__				=			{
         "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,
     }

     A__				=			UNetaDConditionModel(**lowercase_							)
     return model
 @property
 def      UpperCAmelCase_							(				self      :Any							)->							Optional[Any]:
     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      :Dict							)->							Union[str, Any]:
     torch.manual_seed(0							)
     A__				=			VQModel(**self.dummy_movq_kwargs							)
     return model
 def      UpperCAmelCase_							(				self      :Dict							)->							Optional[Any]:
     A__				=			self.dummy_unet
     A__				=			self.dummy_movq

     A__				=			DDIMScheduler(
         num_train_timesteps=10_00						,       beta_schedule="linear"						,       beta_start=0.0_0_0_8_5						,       beta_end=0.0_1_2						,       clip_sample=lowercase_						,       set_alpha_to_one=lowercase_						,       steps_offset=1						,       prediction_type="epsilon"						,       thresholding=lowercase_						,       )

     A__				=			{
         "unet": unet,
         "scheduler": scheduler,
         "movq": movq,
     }
     return components
 def      UpperCAmelCase_							(				self      :Dict						,       lowercase_      :int						,       lowercase_      :Union[str, Any]=0							)->							Dict:
     A__				=			floats_tensor((1, self.text_embedder_hidden_size)						,       rng=random.Random(lowercase_							)							).to(lowercase_							)
     A__				=			floats_tensor((1, self.text_embedder_hidden_size)						,       rng=random.Random(seed + 1							)							).to(
         lowercase_							)
     if str(lowercase_							).startswith("mps"							):
         A__				=			torch.manual_seed(lowercase_							)
     else:
         A__				=			torch.Generator(device=lowercase_							).manual_seed(lowercase_							)
     A__				=			{
         "image_embeds": image_embeds,
         "negative_image_embeds": negative_image_embeds,
         "generator": generator,
         "height": 64,
         "width": 64,
         "guidance_scale": 4.0,
         "num_inference_steps": 2,
         "output_type": "np",
     }
     return inputs



 def      UpperCAmelCase_							(				self      :Optional[Any]							)->							str:
     A__				=			"cpu"

     A__				=			self.get_dummy_components()

     A__				=			self.pipeline_class(**lowercase_							)
     A__				=			pipe.to(lowercase_							)

     pipe.set_progress_bar_config(disable=lowercase_							)

     A__				=			pipe(**self.get_dummy_inputs(lowercase_							)							)
     A__				=			output.images

     A__				=			pipe(
         **self.get_dummy_inputs(lowercase_							)						,       return_dict=lowercase_						,       )[0]

     A__				=			image[0, -3:, -3:, -1]
     A__				=			image_from_tuple[0, -3:, -3:, -1]

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

     A__				=			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    UpperCAmelCase	(       unittest.TestCase				):
 def      UpperCAmelCase_							(				self      :Dict							)->							str:
     # clean up the VRAM after each test
     super().tearDown()
     gc.collect()
     torch.cuda.empty_cache()



 def      UpperCAmelCase_							(				self      :Any							)->							List[Any]:
     A__				=			load_numpy(
         "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
         "/kandinskyv22/kandinskyv22_text2img_cat_fp16.npy"							)

     A__				=			KandinskyVaaPriorPipeline.from_pretrained(
         "kandinsky-community/kandinsky-2-2-prior"						,       torch_dtype=torch.floataa							)
     pipe_prior.to(lowercase_							)

     A__				=			KandinskyVaaPipeline.from_pretrained(
         "kandinsky-community/kandinsky-2-2-decoder"						,       torch_dtype=torch.floataa							)
     A__				=			pipeline.to(lowercase_							)
     pipeline.set_progress_bar_config(disable=lowercase_							)

     A__				=			"red cat, 4k photo"

     A__				=			torch.Generator(device="cuda"							).manual_seed(0							)
     A__,	A__				=			pipe_prior(
         lowercase_						,       generator=lowercase_						,       num_inference_steps=5						,       negative_prompt=""						,       ).to_tuple()

     A__				=			torch.Generator(device="cuda"							).manual_seed(0							)
     A__				=			pipeline(
         image_embeds=lowercase_						,       negative_image_embeds=lowercase_						,       generator=lowercase_						,       num_inference_steps=1_00						,       output_type="np"						,       )

     A__				=			output.images[0]

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

     assert_mean_pixel_difference(lowercase_						,       lowercase_							)



