code: 
  
  
  
  """simple docstring"""
  
  
  
  
  import random
  import unittest
  
  import numpy as np
  import torch
  from PIL import Image
  from torch import nn
  from transformers import (
      CLIPImageProcessor,
      CLIPTextConfig,
      CLIPTextModelWithProjection,
      CLIPTokenizer,
      CLIPVisionConfig,
      CLIPVisionModelWithProjection,
  )
  
  from diffusers import KandinskyVaaPriorEmbaEmbPipeline, PriorTransformer, UnCLIPScheduler
  from diffusers.utils import floats_tensor, torch_device
  from diffusers.utils.testing_utils import enable_full_determinism, skip_mps
  
  from ..test_pipelines_common import PipelineTesterMixin
  
  
  enable_full_determinism()
  
  
  
  
  
  
  class 			UpperCAmelCase_						(       _UpperCamelCase ,    unittest.TestCase		):
        __SCREAMING_SNAKE_CASE		:     int				  =				KandinskyVaaPriorEmbaEmbPipeline
        __SCREAMING_SNAKE_CASE		:     Optional[int]				  =				['prompt', 'image']
        __SCREAMING_SNAKE_CASE		:     Any				  =				['prompt', 'image']
        __SCREAMING_SNAKE_CASE		:     List[Any]				  =				[
            'num_images_per_prompt',
            'strength',
            'generator',
            'num_inference_steps',
            'latents',
            'negative_prompt',
            'guidance_scale',
            'output_type',
            'return_dict',
        ]
        __SCREAMING_SNAKE_CASE		:     Tuple				  =				False
  
  
  
        @property
        def 		snake_case_      (  self  :			Union[str, Any]		):
             return 3_2
  
  
  
        @property
        def 		snake_case_      (  self  :			List[str]		):
             return 3_2
  
  
  
        @property
        def 		snake_case_      (  self  :			Union[str, Any]		):
             return self.time_input_dim
  
  
  
        @property
        def 		snake_case_      (  self  :			int		):
             return self.time_input_dim * 4
  
  
  
        @property
        def 		snake_case_      (  self  :			Dict		):
             return 1_0_0
  
  
  
        @property
        def 		snake_case_      (  self  :			str		):
             _UpperCAmelCase    :			Tuple              =					CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"		)
             return tokenizer
  
  
  
        @property
        def 		snake_case_      (  self  :			Any		):
             torch.manual_seed(0		)
             _UpperCAmelCase    :			Dict              =					CLIPTextConfig(
                 bos_token_id=0  ,			eos_token_id=2  ,			hidden_size=self.text_embedder_hidden_size  ,			projection_dim=self.text_embedder_hidden_size  ,			intermediate_size=3_7  ,			layer_norm_eps=1e-05  ,			num_attention_heads=4  ,			num_hidden_layers=5  ,			pad_token_id=1  ,			vocab_size=1_0_0_0  ,			)
             return CLIPTextModelWithProjection(A		)
  
  
  
        @property
        def 		snake_case_      (  self  :			Optional[Any]		):
             torch.manual_seed(0		)
  
             _UpperCAmelCase    :			List[str]              =					{
                 "num_attention_heads": 2,
                 "attention_head_dim": 1_2,
                 "embedding_dim": self.text_embedder_hidden_size,
                 "num_layers": 1,
             }
  
             _UpperCAmelCase    :			str              =					PriorTransformer(**A		)
             # clip_std and clip_mean is initialized to be 0 so PriorTransformer.post_process_latents will always return 0 - set clip_std to be 1 so it won't return 0
             _UpperCAmelCase    :			str              =					nn.Parameter(torch.ones(model.clip_std.shape		)		)
             return model
  
  
  
        @property
        def 		snake_case_      (  self  :			int		):
             torch.manual_seed(0		)
             _UpperCAmelCase    :			Optional[int]              =					CLIPVisionConfig(
                 hidden_size=self.text_embedder_hidden_size  ,			image_size=2_2_4  ,			projection_dim=self.text_embedder_hidden_size  ,			intermediate_size=3_7  ,			num_attention_heads=4  ,			num_channels=3  ,			num_hidden_layers=5  ,			patch_size=1_4  ,			)
  
             _UpperCAmelCase    :			str              =					CLIPVisionModelWithProjection(A		)
             return model
  
  
  
        @property
        def 		snake_case_      (  self  :			Optional[Any]		):
             _UpperCAmelCase    :			List[Any]              =					CLIPImageProcessor(
                 crop_size=2_2_4  ,			do_center_crop=A  ,			do_normalize=A  ,			do_resize=A  ,			image_mean=[0.48_145_466, 0.4_578_275, 0.40_821_073]  ,			image_std=[0.26_862_954, 0.26_130_258, 0.27_577_711]  ,			resample=3  ,			size=2_2_4  ,			)
  
             return image_processor
  
  
  
        def 		snake_case_      (  self  :			List[Any]		):
             _UpperCAmelCase    :			Any              =					self.dummy_prior
             _UpperCAmelCase    :			Any              =					self.dummy_image_encoder
             _UpperCAmelCase    :			str              =					self.dummy_text_encoder
             _UpperCAmelCase    :			List[str]              =					self.dummy_tokenizer
             _UpperCAmelCase    :			List[Any]              =					self.dummy_image_processor
  
             _UpperCAmelCase    :			str              =					UnCLIPScheduler(
                 variance_type="fixed_small_log"  ,			prediction_type="sample"  ,			num_train_timesteps=1_0_0_0  ,			clip_sample=A  ,			clip_sample_range=10.0  ,			)
  
             _UpperCAmelCase    :			int              =					{
                 "prior": prior,
                 "image_encoder": image_encoder,
                 "text_encoder": text_encoder,
                 "tokenizer": tokenizer,
                 "scheduler": scheduler,
                 "image_processor": image_processor,
             }
  
             return components
  
  
  
        def 		snake_case_      (  self  :			List[Any]  ,			A  :			Tuple  ,			A  :			Tuple=0		):
             if str(A		).startswith("mps"		):
                  _UpperCAmelCase    :			int              =					torch.manual_seed(A		)
             else:
                  _UpperCAmelCase    :			Optional[Any]              =					torch.Generator(device=A		).manual_seed(A		)
  
             _UpperCAmelCase    :			Optional[int]              =					floats_tensor((1, 3, 6_4, 6_4)  ,			rng=random.Random(A		)		).to(A		)
             _UpperCAmelCase    :			Any              =					image.cpu().permute(0  ,			2  ,			3  ,			1		)[0]
             _UpperCAmelCase    :			Union[str, Any]              =					Image.fromarray(np.uinta(A		)		).convert("RGB"		).resize((2_5_6, 2_5_6)		)
  
             _UpperCAmelCase    :			Any              =					{
                 "prompt": "horse",
                 "image": init_image,
                 "strength": 0.5,
                 "generator": generator,
                 "guidance_scale": 4.0,
                 "num_inference_steps": 2,
                 "output_type": "np",
             }
             return inputs
  
  
  
        def 		snake_case_      (  self  :			List[str]		):
             _UpperCAmelCase    :			Any              =					"cpu"
  
             _UpperCAmelCase    :			Dict              =					self.get_dummy_components()
  
             _UpperCAmelCase    :			List[Any]              =					self.pipeline_class(**A		)
             _UpperCAmelCase    :			List[str]              =					pipe.to(A		)
  
             pipe.set_progress_bar_config(disable=A		)
  
             _UpperCAmelCase    :			int              =					pipe(**self.get_dummy_inputs(A		)		)
             _UpperCAmelCase    :			Any              =					output.image_embeds
  
             _UpperCAmelCase    :			Optional[Any]              =					pipe(
                 **self.get_dummy_inputs(A		)  ,			return_dict=A  ,			)[0]
  
             _UpperCAmelCase    :			str              =					image[0, -1_0:]
             _UpperCAmelCase    :			List[str]              =					image_from_tuple[0, -1_0:]
  
             assert image.shape == (1, 3_2)
  
             _UpperCAmelCase    :			Tuple              =					np.array(
                 [
                     0.1_071_284,
                     1.3_330_271,
                     0.61_260_223,
                     -0.6_691_065,
                     -0.3_846_852,
                     -1.0_303_661,
                     0.22_716_111,
                     0.03_348_901,
                     0.30_040_675,
                     -0.24_805_029,
                 ]		)
  
             assert np.abs(image_slice.flatten() - expected_slice		).max() < 1e-2
             assert np.abs(image_from_tuple_slice.flatten() - expected_slice		).max() < 1e-2
  
  
  
        @skip_mps
        def 		snake_case_      (  self  :			List[str]		):
             _UpperCAmelCase    :			str              =					torch_device == "cpu"
             _UpperCAmelCase    :			List[str]              =					True
             _UpperCAmelCase    :			List[str]              =					False
  
             self._test_inference_batch_single_identical(
                 test_max_difference=A  ,			relax_max_difference=A  ,			test_mean_pixel_difference=A  ,			)
  
  
  
        @skip_mps
        def 		snake_case_      (  self  :			List[str]		):
             _UpperCAmelCase    :			Dict              =					torch_device == "cpu"
             _UpperCAmelCase    :			Dict              =					False
  
             self._test_attention_slicing_forward_pass(
                 test_max_difference=A  ,			test_mean_pixel_difference=A  ,			)
  
style_context: 
  
  
  
  """simple docstring"""
  
  
  
  
  import os
  import sys
  from pathlib import Path
  
  import torch
  
  from transformers.testing_utils import TestCasePlus, execute_subprocess_async, require_torch_multi_gpu
  from utils import load_json
  
  
  _lowerCAmelCase       :      Optional[Any]							=							torch.cuda.is_available()
  _lowerCAmelCase       :      str							=							[" Sam ate lunch today.", "Sams lunch ingredients."]
  _lowerCAmelCase       :      Dict							=							["A very interesting story about what I ate for lunch.", "Avocado, celery, turkey, coffee"]
  _lowerCAmelCase       :      Union[str, Any]							=							{
      "max_tokens_per_batch": None,
      "supervise_forward": True,
      "normalize_hidden": True,
      "label_smoothing": 0.2,
      "eval_max_gen_length": None,
      "eval_beams": 1,
      "val_metric": "loss",
      "save_top_k": 1,
      "adafactor": True,
      "early_stopping_patience": 2,
      "logger_name": "default",
      "length_penalty": 0.5,
      "cache_dir": "",
      "task": "summarization",
      "num_workers": 2,
      "alpha_hid": 0,
      "freeze_embeds": True,
      "enc_only": False,
      "tgt_suffix": "",
      "resume_from_checkpoint": None,
      "sortish_sampler": True,
      "student_decoder_layers": 1,
      "val_check_interval": 1.0,
      "output_dir": "",
      "fp16": False,  # TODO(SS): set this to CUDA_AVAILABLE if ci installs apex or start using native amp
      "no_teacher": False,
      "fp16_opt_level": "O1",
      "gpus": 1 if CUDA_AVAILABLE else 0,
      "n_tpu_cores": 0,
      "max_grad_norm": 1.0,
      "do_train": True,
      "do_predict": True,
      "accumulate_grad_batches": 1,
      "server_ip": "",
      "server_port": "",
      "seed": 42,
      "model_name_or_path": "sshleifer/bart-tiny-random",
      "config_name": "",
      "tokenizer_name": "facebook/bart-large",
      "do_lower_case": False,
      "learning_rate": 0.3,
      "lr_scheduler": "linear",
      "weight_decay": 0.0,
      "adam_epsilon": 1E-08,
      "warmup_steps": 0,
      "max_epochs": 1,
      "train_batch_size": 2,
      "eval_batch_size": 2,
      "max_source_length": 12,
      "max_target_length": 12,
      "val_max_target_length": 12,
      "test_max_target_length": 12,
      "fast_dev_run": False,
      "no_cache": False,
      "n_train": -1,
      "n_val": -1,
      "n_test": -1,
      "student_encoder_layers": 1,
      "freeze_encoder": False,
      "auto_scale_batch_size": False,
      "overwrite_output_dir": False,
      "student": None,
  }
  
  def 		__snake_case		(   SCREAMING_SNAKE_CASE__			:		Path		,						SCREAMING_SNAKE_CASE__			:		list       )						->     Dict:
  
  
  
  
  
       '''simple docstring'''
  
       _UpperCAmelCase    :			Tuple              =					"\n".join(SCREAMING_SNAKE_CASE__       )
       Path(SCREAMING_SNAKE_CASE__       ).open("w"       ).writelines(SCREAMING_SNAKE_CASE__       )
  
  def 		__snake_case		(   SCREAMING_SNAKE_CASE__			:		List[Any]       )						->     Union[str, Any]:
  
  
  
  
  
       '''simple docstring'''
  
       for split in ["train", "val", "test"]:
            _dump_articles(os.path.join(SCREAMING_SNAKE_CASE__		,						f'{split}.source'       )		,						SCREAMING_SNAKE_CASE__       )
            _dump_articles(os.path.join(SCREAMING_SNAKE_CASE__		,						f'{split}.target'       )		,						SCREAMING_SNAKE_CASE__       )
       return tmp_dir
  
  
  
  
  
  
  class 			UpperCAmelCase_						(       _UpperCamelCase		):
  
  
  
        @classmethod
        def 		snake_case_      (  cls  :			Optional[int]		):
             return cls
  
  
  
        @require_torch_multi_gpu
        def 		snake_case_      (  self  :			Union[str, Any]		):
             _UpperCAmelCase    :			Dict              =					{
                 "no_teacher": True,
                 "freeze_encoder": True,
                 "gpus": 2,
                 "overwrite_output_dir": True,
                 "sortish_sampler": True,
             }
             self._test_distiller_cli_fork(A  ,			check_contents=A		)
  
  
  
        def 		snake_case_      (  self  :			List[str]  ,			A  :			List[Any]  ,			A  :			Optional[Any]=True		):
             _UpperCAmelCase    :			int              =					{
                 "label_smoothing": 0.0,
                 "early_stopping_patience": -1,
                 "train_batch_size": 1,
                 "eval_batch_size": 2,
                 "max_epochs": 2,
                 "alpha_mlm": 0.2,
                 "alpha_ce": 0.8,
                 "do_predict": True,
                 "model_name_or_path": "sshleifer/tinier_bart",
                 "teacher": CHEAP_ARGS["model_name_or_path"],
                 "val_check_interval": 0.5,
             }
             default_updates.update(A		)
             _UpperCAmelCase    :			dict              =					CHEAP_ARGS.copy()
             _UpperCAmelCase    :			Union[str, Any]              =					make_test_data_dir(tmp_dir=self.get_auto_remove_tmp_dir()		)
             _UpperCAmelCase    :			List[Any]              =					self.get_auto_remove_tmp_dir()
             args_d.update(data_dir=A  ,			output_dir=A  ,			**A		)
  
             def convert(A  :			List[Any]  ,			A  :			Optional[int]		):
                  if k in ["tgt_suffix", "server_ip", "server_port", "out", "n_tpu_cores"]:
                       return ""
                  if v is False or v is None:
                       return ""
                  if v is True:  # or len(str(v))==0:
                       return f'--{k}'
                  return f'--{k}={v}'
  
             _UpperCAmelCase    :			Tuple              =					[x for x in (convert(A  ,			A		) for k, v in args_d.items()) if len(A		)]
             _UpperCAmelCase    :			Union[str, Any]              =					[sys.executable, f'{self.test_file_dir}/distillation.py'] + cli_args
             execute_subprocess_async(A  ,			env=self.get_env()		)
  
             _UpperCAmelCase    :			int              =					os.listdir(A		)
             _UpperCAmelCase    :			List[Any]              =					{os.path.basename(A		) for p in contents}
             _UpperCAmelCase    :			List[Any]              =					[p for p in contents if p.endswith("ckpt"		)]
             assert len(A		) > 0
  
             self.assertIn("test_generations.txt"  ,			A		)
             self.assertIn("test_results.txt"  ,			A		)
  
             # get the following from the module, (we don't have access to `model` here)
             _UpperCAmelCase    :			List[Any]              =					os.path.join(A  ,			"metrics.json"		)
             _UpperCAmelCase    :			str              =					"rouge2"
  
             _UpperCAmelCase    :			List[Any]              =					load_json(A		)
             # {'test': [{'test_avg_loss': 10.63731575012207, 'test_avg_rouge1': 0.0, 'test_avg_rouge2': 0.0, 'test_avg_rougeL': 0.0, 'test_avg_gen_time': 0.1822289228439331, 'test_avg_gen_len': 142.0, 'step_count': 1}]}
             print(A		)
             _UpperCAmelCase    :			List[str]              =					metrics["val"][-1]
             self.assertGreaterEqual(last_step_stats["val_avg_gen_time"]  ,			0.01		)
             self.assertIsInstance(last_step_stats[f'val_avg_{val_metric}']  ,			A		)
             self.assertEqual(len(metrics["test"]		)  ,			1		)
             _UpperCAmelCase    :			Dict              =					int(args_d["max_epochs"] * (1 / args_d["val_check_interval"]) / 2 + 1		)
             self.assertEqual(len(metrics["val"]		)  ,			A		)
  