import argparse
import gc
import json
import os

import torch
from datasets import load_dataset
from torch.optim import AdamW
from torch.utils.data import DataLoader
from transformers import AutoModelForSequenceClassification, AutoTokenizer, get_linear_schedule_with_warmup, set_seed

from accelerate import Accelerator, DistributedType
from accelerate.utils.deepspeed import DummyOptim, DummyScheduler


UpperCAmelCase_								=							16
UpperCAmelCase_								=							32

def 							UpperCamelCase				(    lowerCAmelCase_   )		->						int:




       '''simple docstring'''






       return int(x / 2**20   )





class      lowerCAmelCase						:



     def __enter__(					self       ):
            gc.collect()
            torch.cuda.empty_cache()
            torch.cuda.reset_max_memory_allocated()  # reset the peak gauge to zero
            _A=	torch.cuda.memory_allocated()
            return self





     def __exit__(					self				,      *lowerCAmelCase__       ):
            gc.collect()
            torch.cuda.empty_cache()
            _A=	torch.cuda.memory_allocated()
            _A=	torch.cuda.max_memory_allocated()
            _A=	bamb(self.end - self.begin       )
            _A=	bamb(self.peak - self.begin       )
            # print(f"delta used/peak {self.used:4d}/{self.peaked:4d}")


def 							UpperCamelCase				(    lowerCAmelCase_       ,					lowerCAmelCase_ = 16       ,					lowerCAmelCase_ = "bert-base-cased"       ,					lowerCAmelCase_ = 3_20       ,					lowerCAmelCase_ = 1_60       ,					)		->						Optional[Any]:




       '''simple docstring'''






       _A=	AutoTokenizer.from_pretrained(lowerCAmelCase_   )
       _A=	load_dataset(
           'glue'       ,					'mrpc'       ,					split={'train': F"train[:{n_train}]", 'validation': F"validation[:{n_val}]"}   )

       def tokenize_function(lowerCAmelCase_   ):
              # max_length=None => use the model max length (it's actually the default)
              _A=	tokenizer(examples['sentence1']       ,					examples['sentence2']       ,					truncation=lowerCAmelCase_       ,					max_length=lowerCAmelCase_   )
              return outputs

       # Apply the method we just defined to all the examples in all the splits of the dataset
       _A=	datasets.map(
           lowerCAmelCase_       ,					batched=lowerCAmelCase_       ,					remove_columns=['idx', 'sentence1', 'sentence2']       ,					load_from_cache_file=lowerCAmelCase_   )

       # We also rename the 'label' column to 'labels' which is the expected name for labels by the models of the
       # transformers library
       _A=	tokenized_datasets.rename_column('label'       ,					'labels'   )

       def collate_fn(lowerCAmelCase_   ):
              # On TPU it's best to pad everything to the same length or training will be very slow.
              if accelerator.distributed_type == DistributedType.TPU:
                     return tokenizer.pad(lowerCAmelCase_       ,					padding='max_length'       ,					max_length=1_28       ,					return_tensors='pt'   )
              return tokenizer.pad(lowerCAmelCase_       ,					padding='longest'       ,					return_tensors='pt'   )

       # Instantiate dataloaders.
       _A=	DataLoader(
           tokenized_datasets['train']       ,					shuffle=lowerCAmelCase_       ,					collate_fn=lowerCAmelCase_       ,					batch_size=lowerCAmelCase_   )
       _A=	DataLoader(
           tokenized_datasets['validation']       ,					shuffle=lowerCAmelCase_       ,					collate_fn=lowerCAmelCase_       ,					batch_size=lowerCAmelCase_   )

       return train_dataloader, eval_dataloader

def 							UpperCamelCase				(    lowerCAmelCase_       ,					lowerCAmelCase_   )		->						Optional[int]:




       '''simple docstring'''






       _A=	Accelerator()

       # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
       _A=	config['lr']
       _A=	int(config['num_epochs']   )
       _A=	int(config['seed']   )
       _A=	int(config['batch_size']   )
       _A=	args.model_name_or_path

       set_seed(lowerCAmelCase_   )
       _A,			_A=	get_dataloaders(lowerCAmelCase_       ,					lowerCAmelCase_       ,					lowerCAmelCase_       ,					args.n_train       ,					args.n_val   )

       # Instantiate the model (we build the model here so that the seed also control new weights initialization)
       _A=	AutoModelForSequenceClassification.from_pretrained(lowerCAmelCase_       ,					return_dict=lowerCAmelCase_   )

       # Instantiate optimizer
       _A=	(
           AdamW
           if accelerator.state.deepspeed_plugin is None
           or 'optimizer' not in accelerator.state.deepspeed_plugin.deepspeed_config
           else DummyOptim
       )
       _A=	optimizer_cls(params=model.parameters()       ,					lr=lowerCAmelCase_   )

       if accelerator.state.deepspeed_plugin is not None:
              _A=	accelerator.state.deepspeed_plugin.deepspeed_config[
                  'gradient_accumulation_steps'
              ]
       else:
              _A=	1
       _A=	(len(lowerCAmelCase_   ) * num_epochs) // gradient_accumulation_steps

       # Instantiate scheduler
       if (
           accelerator.state.deepspeed_plugin is None
           or "scheduler" not in accelerator.state.deepspeed_plugin.deepspeed_config
       ):
              _A=	get_linear_schedule_with_warmup(
                  optimizer=lowerCAmelCase_       ,					num_warmup_steps=0       ,					num_training_steps=lowerCAmelCase_       ,					)
       else:
              _A=	DummyScheduler(lowerCAmelCase_       ,					total_num_steps=lowerCAmelCase_       ,					warmup_num_steps=0   )

       # Prepare everything
       # There is no specific order to remember, we just need to unpack the objects in the same order we gave them to the
       # prepare method.
       _A,			_A,			_A,			_A,			_A=	accelerator.prepare(
           lowerCAmelCase_       ,					lowerCAmelCase_       ,					lowerCAmelCase_       ,					lowerCAmelCase_       ,					lowerCAmelCase_   )

       # We need to keep track of how many total steps we have iterated over
       _A=	0
       # We also need to keep track of the stating epoch so files are named properly
       _A=	0

       # Now we train the model
       _A=	{}
       for epoch in range(lowerCAmelCase_       ,					lowerCAmelCase_   ):
              with TorchTracemalloc() as tracemalloc:
                     model.train()
                     for step, batch in enumerate(lowerCAmelCase_   ):
                            _A=	model(**lowerCAmelCase_   )
                            _A=	outputs.loss
                            _A=	loss / gradient_accumulation_steps
                            accelerator.backward(lowerCAmelCase_   )
                            if step % gradient_accumulation_steps == 0:
                                   optimizer.step()
                                   lr_scheduler.step()
                                   optimizer.zero_grad()

                            overall_step += 1

        # Printing the GPU memory usage details such as allocated memory, peak memory, and total memory usage
              accelerator.print('Memory before entering the train : {}'.format(bamb(tracemalloc.begin   )   )   )
              accelerator.print('Memory consumed at the end of the train (end-begin): {}'.format(tracemalloc.used   )   )
              accelerator.print('Peak Memory consumed during the train (max-begin): {}'.format(tracemalloc.peaked   )   )
              accelerator.print(
                  'Total Peak Memory consumed during the train (max): {}'.format(
                      tracemalloc.peaked + bamb(tracemalloc.begin   )   )   )
              _A=	tracemalloc.peaked + bamb(tracemalloc.begin   )
              if args.peak_memory_upper_bound is not None:
                     assert (
                         train_total_peak_memory[F"epoch-{epoch}"] <= args.peak_memory_upper_bound
                     ), "Peak memory usage exceeded the upper bound"

       accelerator.wait_for_everyone()
       if accelerator.is_main_process:
              with open(os.path.join(args.output_dir       ,					'peak_memory_utilization.json'   )       ,					'w'   ) as f:
                     json.dump(lowerCAmelCase_       ,					lowerCAmelCase_   )

def 							UpperCamelCase				(    )		->						Any:




       '''simple docstring'''






       _A=	argparse.ArgumentParser(description='Simple example of training script tracking peak GPU memory usage.'   )
       parser.add_argument(
           '--model_name_or_path'       ,					type=lowerCAmelCase_       ,					default='bert-base-cased'       ,					help='Path to pretrained model or model identifier from huggingface.co/models.'       ,					required=lowerCAmelCase_       ,					)
       parser.add_argument(
           '--output_dir'       ,					type=lowerCAmelCase_       ,					default='.'       ,					help='Optional save directory where all checkpoint folders will be stored. Default is the current working directory.'       ,					)
       parser.add_argument(
           '--peak_memory_upper_bound'       ,					type=lowerCAmelCase_       ,					default=lowerCAmelCase_       ,					help='The upper bound of peak memory usage in MB. If set, the training will throw an error if the peak memory usage exceeds this value.'       ,					)
       parser.add_argument(
           '--n_train'       ,					type=lowerCAmelCase_       ,					default=3_20       ,					help='Number of training examples to use.'       ,					)
       parser.add_argument(
           '--n_val'       ,					type=lowerCAmelCase_       ,					default=1_60       ,					help='Number of validation examples to use.'       ,					)
       parser.add_argument(
           '--num_epochs'       ,					type=lowerCAmelCase_       ,					default=1       ,					help='Number of train epochs.'       ,					)
       _A=	parser.parse_args()
       _A=	{'lr': 2E-5, 'num_epochs': args.num_epochs, 'seed': 42, 'batch_size': 16}
       training_function(lowerCAmelCase_       ,					lowerCAmelCase_   )


if __name__ == "__main__":
      main()
import gc
import random
import unittest

import numpy as np
import torch
from transformers import (
    CLIPImageProcessor,
    CLIPTextConfig,
    CLIPTextModelWithProjection,
    CLIPTokenizer,
    CLIPVisionConfig,
    CLIPVisionModelWithProjection,
)

from diffusers import (
    DiffusionPipeline,
    UnCLIPImageVariationPipeline,
    UnCLIPScheduler,
    UNetaDConditionModel,
    UNetaDModel,
)
from diffusers.pipelines.unclip.text_proj import UnCLIPTextProjModel
from diffusers.utils import floats_tensor, load_numpy, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, load_image, require_torch_gpu, skip_mps

from ..pipeline_params import IMAGE_VARIATION_BATCH_PARAMS, IMAGE_VARIATION_PARAMS
from ..test_pipelines_common import PipelineTesterMixin, assert_mean_pixel_difference


enable_full_determinism()





class      lowerCAmelCase						(					_a			,		unittest.TestCase			):
     _SCREAMING_SNAKE_CASE     :     Union[str, Any]		=UnCLIPImageVariationPipeline
     _SCREAMING_SNAKE_CASE     :     Union[str, Any]		=IMAGE_VARIATION_PARAMS - {"""height""", """width""", """guidance_scale"""}
     _SCREAMING_SNAKE_CASE     :     Union[str, Any]		=IMAGE_VARIATION_BATCH_PARAMS

     _SCREAMING_SNAKE_CASE     :     int		=[
         """generator""",
         """return_dict""",
         """decoder_num_inference_steps""",
         """super_res_num_inference_steps""",
     ]
     _SCREAMING_SNAKE_CASE     :     Optional[int]		=False



     @property
     def 			a__  (					self       ):
            return 32



     @property
     def 			a__  (					self       ):
            return 32



     @property
     def 			a__  (					self       ):
            return self.time_input_dim



     @property
     def 			a__  (					self       ):
            return self.time_input_dim * 4



     @property
     def 			a__  (					self       ):
            return 100



     @property
     def 			a__  (					self       ):
            _A=	CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'       )
            return tokenizer



     @property
     def 			a__  (					self       ):
            torch.manual_seed(0       )
            _A=	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=37				,      layer_norm_eps=1E-05				,      num_attention_heads=4				,      num_hidden_layers=5				,      pad_token_id=1				,      vocab_size=1000				,      )
            return CLIPTextModelWithProjection(lowerCAmelCase__       )



     @property
     def 			a__  (					self       ):
            torch.manual_seed(0       )
            _A=	CLIPVisionConfig(
                hidden_size=self.text_embedder_hidden_size				,      projection_dim=self.text_embedder_hidden_size				,      num_hidden_layers=5				,      num_attention_heads=4				,      image_size=32				,      intermediate_size=37				,      patch_size=1				,      )
            return CLIPVisionModelWithProjection(lowerCAmelCase__       )



     @property
     def 			a__  (					self       ):
            torch.manual_seed(0       )

            _A=	{
                'clip_embeddings_dim': self.text_embedder_hidden_size,
                'time_embed_dim': self.time_embed_dim,
                'cross_attention_dim': self.cross_attention_dim,
            }

            _A=	UnCLIPTextProjModel(**lowerCAmelCase__       )
            return model



     @property
     def 			a__  (					self       ):
            torch.manual_seed(0       )

            _A=	{
                'sample_size': 32,
                # RGB in channels
                'in_channels': 3,
                # Out channels is double in channels because predicts mean and variance
                'out_channels': 6,
                '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,
                'cross_attention_dim': self.cross_attention_dim,
                'attention_head_dim': 4,
                'resnet_time_scale_shift': 'scale_shift',
                'class_embed_type': 'identity',
            }

            _A=	UNetaDConditionModel(**lowerCAmelCase__       )
            return model



     @property
     def 			a__  (					self       ):
            return {
                "sample_size": 64,
                "layers_per_block": 1,
                "down_block_types": ("ResnetDownsampleBlock2D", "ResnetDownsampleBlock2D"),
                "up_block_types": ("ResnetUpsampleBlock2D", "ResnetUpsampleBlock2D"),
                "block_out_channels": (self.block_out_channels_a, self.block_out_channels_a * 2),
                "in_channels": 6,
                "out_channels": 3,
            }



     @property
     def 			a__  (					self       ):
            torch.manual_seed(0       )

            _A=	UNetaDModel(**self.dummy_super_res_kwargs       )
            return model



     @property
     def 			a__  (					self       ):
            # seeded differently to get different unet than `self.dummy_super_res_first`
            torch.manual_seed(1       )

            _A=	UNetaDModel(**self.dummy_super_res_kwargs       )
            return model



     def 			a__  (					self       ):
            _A=	self.dummy_decoder
            _A=	self.dummy_text_proj
            _A=	self.dummy_text_encoder
            _A=	self.dummy_tokenizer
            _A=	self.dummy_super_res_first
            _A=	self.dummy_super_res_last

            _A=	UnCLIPScheduler(
                variance_type='learned_range'				,      prediction_type='epsilon'				,      num_train_timesteps=1000				,      )

            _A=	UnCLIPScheduler(
                variance_type='fixed_small_log'				,      prediction_type='epsilon'				,      num_train_timesteps=1000				,      )

            _A=	CLIPImageProcessor(crop_size=32				,      size=32       )

            _A=	self.dummy_image_encoder

            return {
                "decoder": decoder,
                "text_encoder": text_encoder,
                "tokenizer": tokenizer,
                "text_proj": text_proj,
                "feature_extractor": feature_extractor,
                "image_encoder": image_encoder,
                "super_res_first": super_res_first,
                "super_res_last": super_res_last,
                "decoder_scheduler": decoder_scheduler,
                "super_res_scheduler": super_res_scheduler,
            }



     def 			a__  (					self				,      lowerCAmelCase__				,      lowerCAmelCase__=0				,      lowerCAmelCase__=True       ):
            _A=	floats_tensor((1, 3, 32, 32)				,      rng=random.Random(lowerCAmelCase__       )       ).to(lowerCAmelCase__       )
            if str(lowerCAmelCase__       ).startswith('mps'       ):
                   _A=	torch.manual_seed(lowerCAmelCase__       )
            else:
                   _A=	torch.Generator(device=lowerCAmelCase__       ).manual_seed(lowerCAmelCase__       )

            if pil_image:
                   _A=	input_image * 0.5 + 0.5
                   _A=	input_image.clamp(0				,      1       )
                   _A=	input_image.cpu().permute(0				,      2				,      3				,      1       ).float().numpy()
                   _A=	DiffusionPipeline.numpy_to_pil(lowerCAmelCase__       )[0]

            return {
                "image": input_image,
                "generator": generator,
                "decoder_num_inference_steps": 2,
                "super_res_num_inference_steps": 2,
                "output_type": "np",
            }



     def 			a__  (					self       ):
            _A=	'cpu'

            _A=	self.get_dummy_components()

            _A=	self.pipeline_class(**lowerCAmelCase__       )
            _A=	pipe.to(lowerCAmelCase__       )

            pipe.set_progress_bar_config(disable=lowerCAmelCase__       )

            _A=	self.get_dummy_inputs(lowerCAmelCase__				,      pil_image=lowerCAmelCase__       )

            _A=	pipe(**lowerCAmelCase__       )
            _A=	output.images

            _A=	self.get_dummy_inputs(lowerCAmelCase__				,      pil_image=lowerCAmelCase__       )

            _A=	pipe(
                **lowerCAmelCase__				,      return_dict=lowerCAmelCase__				,      )[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.9997,
                    0.0002,
                    0.9997,
                    0.9997,
                    0.9969,
                    0.0023,
                    0.9997,
                    0.9969,
                    0.9970,
                ]       )

            assert np.abs(image_slice.flatten() - expected_slice       ).max() < 1E-2
            assert np.abs(image_from_tuple_slice.flatten() - expected_slice       ).max() < 1E-2



     def 			a__  (					self       ):
            _A=	'cpu'

            _A=	self.get_dummy_components()

            _A=	self.pipeline_class(**lowerCAmelCase__       )
            _A=	pipe.to(lowerCAmelCase__       )

            pipe.set_progress_bar_config(disable=lowerCAmelCase__       )

            _A=	self.get_dummy_inputs(lowerCAmelCase__				,      pil_image=lowerCAmelCase__       )

            _A=	pipe(**lowerCAmelCase__       )
            _A=	output.images

            _A=	self.get_dummy_inputs(lowerCAmelCase__				,      pil_image=lowerCAmelCase__       )

            _A=	pipe(
                **lowerCAmelCase__				,      return_dict=lowerCAmelCase__				,      )[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.9997, 0.0003, 0.9997, 0.9997, 0.9970, 0.0024, 0.9997, 0.9971, 0.9971]       )

            assert np.abs(image_slice.flatten() - expected_slice       ).max() < 1E-2
            assert np.abs(image_from_tuple_slice.flatten() - expected_slice       ).max() < 1E-2



     def 			a__  (					self       ):
            _A=	'cpu'

            _A=	self.get_dummy_components()

            _A=	self.pipeline_class(**lowerCAmelCase__       )
            _A=	pipe.to(lowerCAmelCase__       )

            pipe.set_progress_bar_config(disable=lowerCAmelCase__       )

            _A=	self.get_dummy_inputs(lowerCAmelCase__				,      pil_image=lowerCAmelCase__       )
            _A=	[
                pipeline_inputs['image'],
                pipeline_inputs['image'],
            ]

            _A=	pipe(**lowerCAmelCase__       )
            _A=	output.images

            _A=	self.get_dummy_inputs(lowerCAmelCase__				,      pil_image=lowerCAmelCase__       )
            _A=	[
                tuple_pipeline_inputs['image'],
                tuple_pipeline_inputs['image'],
            ]

            _A=	pipe(
                **lowerCAmelCase__				,      return_dict=lowerCAmelCase__				,      )[0]

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

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

            _A=	np.array(
                [
                    0.9997,
                    0.9989,
                    0.0008,
                    0.0021,
                    0.9960,
                    0.0018,
                    0.0014,
                    0.0002,
                    0.9933,
                ]       )

            assert np.abs(image_slice.flatten() - expected_slice       ).max() < 1E-2
            assert np.abs(image_from_tuple_slice.flatten() - expected_slice       ).max() < 1E-2



     def 			a__  (					self       ):
            _A=	torch.device('cpu'       )





            class      lowerCAmelCase						:
                 _SCREAMING_SNAKE_CASE     :     List[str]		=1


            _A=	self.get_dummy_components()

            _A=	self.pipeline_class(**lowerCAmelCase__       )
            _A=	pipe.to(lowerCAmelCase__       )

            pipe.set_progress_bar_config(disable=lowerCAmelCase__       )

            _A=	torch.Generator(device=lowerCAmelCase__       ).manual_seed(0       )
            _A=	pipe.decoder.dtype
            _A=	1

            _A=	(
                batch_size,
                pipe.decoder.config.in_channels,
                pipe.decoder.config.sample_size,
                pipe.decoder.config.sample_size,
            )
            _A=	pipe.prepare_latents(
                lowerCAmelCase__				,      dtype=lowerCAmelCase__				,      device=lowerCAmelCase__				,      generator=lowerCAmelCase__				,      latents=lowerCAmelCase__				,      scheduler=DummyScheduler()       )

            _A=	(
                batch_size,
                pipe.super_res_first.config.in_channels // 2,
                pipe.super_res_first.config.sample_size,
                pipe.super_res_first.config.sample_size,
            )
            _A=	pipe.prepare_latents(
                lowerCAmelCase__				,      dtype=lowerCAmelCase__				,      device=lowerCAmelCase__				,      generator=lowerCAmelCase__				,      latents=lowerCAmelCase__				,      scheduler=DummyScheduler()       )

            _A=	self.get_dummy_inputs(lowerCAmelCase__				,      pil_image=lowerCAmelCase__       )

            _A=	pipe(
                **lowerCAmelCase__				,      decoder_latents=lowerCAmelCase__				,      super_res_latents=lowerCAmelCase__       ).images

            _A=	self.get_dummy_inputs(lowerCAmelCase__				,      pil_image=lowerCAmelCase__       )
            # Don't pass image, instead pass embedding
            _A=	pipeline_inputs.pop('image'       )
            _A=	pipe.image_encoder(lowerCAmelCase__       ).image_embeds

            _A=	pipe(
                **lowerCAmelCase__				,      decoder_latents=lowerCAmelCase__				,      super_res_latents=lowerCAmelCase__				,      image_embeddings=lowerCAmelCase__				,      ).images

            # make sure passing text embeddings manually is identical
            assert np.abs(img_out_a - img_out_a       ).max() < 1E-4



     @skip_mps
     def 			a__  (					self       ):
            _A=	torch_device == 'cpu'

            # Check is relaxed because there is not a torch 2.0 sliced attention added kv processor
            _A=	1E-2

            self._test_attention_slicing_forward_pass(
                test_max_difference=lowerCAmelCase__				,      expected_max_diff=lowerCAmelCase__       )



     @skip_mps
     def 			a__  (					self       ):
            _A=	torch_device == 'cpu'
            _A=	True
            _A=	[
                'decoder_num_inference_steps',
                'super_res_num_inference_steps',
            ]

            self._test_inference_batch_single_identical(
                test_max_difference=lowerCAmelCase__				,      relax_max_difference=lowerCAmelCase__				,      additional_params_copy_to_batched_inputs=lowerCAmelCase__				,      )



     def 			a__  (					self       ):
            _A=	[
                'decoder_num_inference_steps',
                'super_res_num_inference_steps',
            ]

            if torch_device == "mps":
                   # TODO: MPS errors with larger batch sizes
                   _A=	[2, 3]
                   self._test_inference_batch_consistent(
                       batch_sizes=lowerCAmelCase__				,      additional_params_copy_to_batched_inputs=lowerCAmelCase__				,      )
            else:
                   self._test_inference_batch_consistent(
                       additional_params_copy_to_batched_inputs=lowerCAmelCase__       )



     @skip_mps
     def 			a__  (					self       ):
            return super().test_dict_tuple_outputs_equivalent()



     @skip_mps
     def 			a__  (					self       ):
            return super().test_save_load_local()





     @skip_mps
     def 			a__  (					self       ):
            return super().test_save_load_optional_components()





@slow
@require_torch_gpu
class      lowerCAmelCase						(					unittest.TestCase			):



     def 			a__  (					self       ):
            # clean up the VRAM after each test
            super().tearDown()
            gc.collect()
            torch.cuda.empty_cache()





     def 			a__  (					self       ):
            _A=	load_image(
                'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/unclip/cat.png'       )
            _A=	load_numpy(
                'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
                '/unclip/karlo_v1_alpha_cat_variation_fp16.npy'       )

            _A=	UnCLIPImageVariationPipeline.from_pretrained(
                'kakaobrain/karlo-v1-alpha-image-variations'				,      torch_dtype=torch.floataa       )
            _A=	pipeline.to(lowerCAmelCase__       )
            pipeline.set_progress_bar_config(disable=lowerCAmelCase__       )

            _A=	torch.Generator(device='cpu'       ).manual_seed(0       )
            _A=	pipeline(
                lowerCAmelCase__				,      generator=lowerCAmelCase__				,      output_type='np'				,      )

            _A=	output.images[0]

            assert image.shape == (256, 256, 3)

            assert_mean_pixel_difference(lowerCAmelCase__				,      lowerCAmelCase__				,      15       )