

import argparse
import os

import evaluate
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


########################################################################
# This is a fully working simple example to use Accelerate,
# specifically showcasing the checkpointing capability,
# and builds off the `nlp_example.py` script.
#
# This example trains a Bert base model on GLUE MRPC
# in any of the following settings (with the same script):
#   - single CPU or single GPU
#   - multi GPUS (using PyTorch distributed mode)
#   - (multi) TPUs
#   - fp16 (mixed-precision) or fp32 (normal precision)
#
# To help focus on the differences in the code, building `DataLoaders`
# was refactored into its own function.
# New additions from the base script can be found quickly by
# looking for the # New Code # tags
#
# To run it in each of these various modes, follow the instructions
# in the readme for examples:
# https://github.com/huggingface/accelerate/tree/main/examples
#
########################################################################

_UpperCAmelCase    	=			16
_UpperCAmelCase    	=			32


def 			__UpperCamelCase    (lowerCAmelCase					:  Accelerator,    lowerCAmelCase					:  int = 16     )     ->  Tuple:
      A    =				AutoTokenizer.from_pretrained('bert-base-cased'     )
      A    =				load_dataset('glue',    'mrpc'     )

      def tokenize_function(lowerCAmelCase					:  Union[str, Any]     ):
            # 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
      # starting with the main process first:
      with accelerator.main_process_first():
            A    =				datasets.map(
                lowerCAmelCase,    batched=lowerCAmelCase,    remove_columns=['idx', 'sentence1', 'sentence2'],    )

      # 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					:  Any     ):
            # On TPU it's best to pad everything to the same length or training will be very slow.
            A    =				128 if accelerator.distributed_type == DistributedType.TPU else None
            # When using mixed precision we want round multiples of 8/16
            if accelerator.mixed_precision == "fp8":
                  A    =				16
            elif accelerator.mixed_precision != "no":
                  A    =				8
            else:
                  A    =				None

            return tokenizer.pad(
                lowerCAmelCase,    padding='longest',    max_length=lowerCAmelCase,    pad_to_multiple_of=lowerCAmelCase,    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


# For testing only
if os.environ.get("TESTING_MOCKED_DATALOADERS", None) == "1":
  from accelerate.test_utils.training import mocked_dataloaders

  _UpperCAmelCase    	=			mocked_dataloaders  # noqa: F811


def 			__UpperCamelCase    (lowerCAmelCase					:  Dict,    lowerCAmelCase					:  Union[str, Any]     )     ->  Optional[Any]:
      # For testing only
      if os.environ.get('TESTING_MOCKED_DATALOADERS',    lowerCAmelCase     ) == "1":
            A    =				2
      # Initialize accelerator
      A    =				Accelerator(cpu=args.cpu,    mixed_precision=args.mixed_precision     )
      # 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']     )

      # New Code #
      # Parse out whether we are saving every epoch or after a certain number of batches
      if hasattr(args.checkpointing_steps,    'isdigit'     ):
            if args.checkpointing_steps == "epoch":
                  A    =				args.checkpointing_steps
            elif args.checkpointing_steps.isdigit():
                  A    =				int(args.checkpointing_steps     )
            else:
                  raise ValueError(
                      f'''Argument `checkpointing_steps` must be either a number or `epoch`. `{args.checkpointing_steps}` passed.'''     )
      else:
            A    =				None

      set_seed(lowerCAmelCase     )

      A ,	A    =				get_dataloaders(lowerCAmelCase,    lowerCAmelCase     )
      A    =				evaluate.load('glue',    'mrpc'     )

      # If the batch size is too big we use gradient accumulation
      A    =				1
      if batch_size > MAX_GPU_BATCH_SIZE and accelerator.distributed_type != DistributedType.TPU:
            A    =				batch_size // MAX_GPU_BATCH_SIZE
            A    =				MAX_GPU_BATCH_SIZE

      # Instantiate the model (we build the model here so that the seed also control new weights initialization)
      A    =				AutoModelForSequenceClassification.from_pretrained('bert-base-cased',    return_dict=lowerCAmelCase     )

      # We could avoid this line since the accelerator is set with `device_placement=True` (default value).
      # Note that if you are placing tensors on devices manually, this line absolutely needs to be before the optimizer
      # creation otherwise training will not work on TPU (`accelerate` will kindly throw an error to make us aware of that).
      A    =				model.to(accelerator.device     )

      # Instantiate optimizer
      A    =				AdamW(params=model.parameters(),    lr=lowerCAmelCase     )

      # Instantiate scheduler
      A    =				get_linear_schedule_with_warmup(
          optimizer=lowerCAmelCase,    num_warmup_steps=100,    num_training_steps=(len(lowerCAmelCase     ) * num_epochs) // gradient_accumulation_steps,    )

      # 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     )

      # New Code #
      # 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

      # We need to load the checkpoint back in before training here with `load_state`
      # The total number of epochs is adjusted based on where the state is being loaded from,
      # as we assume continuation of the same training script
      if args.resume_from_checkpoint:
            if args.resume_from_checkpoint is not None or args.resume_from_checkpoint != "":
                  accelerator.print(f'''Resumed from checkpoint: {args.resume_from_checkpoint}'''     )
                  accelerator.load_state(args.resume_from_checkpoint     )
                  A    =				os.path.basename(args.resume_from_checkpoint     )
            else:
                  # Get the most recent checkpoint
                  A    =				[f.name for f in os.scandir(os.getcwd()     ) if f.is_dir()]
                  dirs.sort(key=os.path.getctime     )
                  A    =				dirs[-1]  # Sorts folders by date modified, most recent checkpoint is the last
            # Extract `epoch_{i}` or `step_{i}`
            A    =				os.path.splitext(lowerCAmelCase     )[0]

            if "epoch" in training_difference:
                  A    =				int(training_difference.replace('epoch_',    ''     )     ) + 1
                  A    =				None
            else:
                  A    =				int(training_difference.replace('step_',    ''     )     )
                  A    =				resume_step // len(lowerCAmelCase     )
                  resume_step -= starting_epoch * len(lowerCAmelCase     )

    # Now we train the model
      for epoch in range(lowerCAmelCase,    lowerCAmelCase     ):
            model.train()
            # New Code #
            if args.resume_from_checkpoint and epoch == starting_epoch and resume_step is not None:
                  # We need to skip steps until we reach the resumed step
                  A    =				accelerator.skip_first_batches(lowerCAmelCase,    lowerCAmelCase     )
                  overall_step += resume_step
            else:
                  # After the first iteration though, we need to go back to the original dataloader
                  A    =				train_dataloader
            for step, batch in enumerate(lowerCAmelCase     ):
                  # We could avoid this line since we set the accelerator with `device_placement=True`.
                  batch.to(accelerator.device     )
                  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()
                  # New Code #
                  overall_step += 1

                  # New Code #
                  # We save the model, optimizer, lr_scheduler, and seed states by calling `save_state`
                  # These are saved to folders named `step_{overall_step}`
                  # Will contain files: "pytorch_model.bin", "optimizer.bin", "scheduler.bin", and "random_states.pkl"
                  # If mixed precision was used, will also save a "scalar.bin" file
                  if isinstance(lowerCAmelCase,    lowerCAmelCase     ):
                        A    =				f'''step_{overall_step}'''
                        if overall_step % checkpointing_steps == 0:
                              if args.output_dir is not None:
                                    A    =				os.path.join(args.output_dir,    lowerCAmelCase     )
                              accelerator.save_state(lowerCAmelCase     )

            model.eval()
            for step, batch in enumerate(lowerCAmelCase     ):
                  # We could avoid this line since we set the accelerator with `device_placement=True` (the default).
                  batch.to(accelerator.device     )
                  with torch.no_grad():
                        A    =				model(**lowerCAmelCase     )
                  A    =				outputs.logits.argmax(dim=-1     )
                  A ,	A    =				accelerator.gather_for_metrics((predictions, batch['labels'])     )
                  metric.add_batch(
                      predictions=lowerCAmelCase,    references=lowerCAmelCase,    )

            A    =				metric.compute()
            # Use accelerator.print to print only on the main process.
            accelerator.print(f'''epoch {epoch}:''',    lowerCAmelCase     )

            # New Code #
            # We save the model, optimizer, lr_scheduler, and seed states by calling `save_state`
            # These are saved to folders named `epoch_{epoch}`
            # Will contain files: "pytorch_model.bin", "optimizer.bin", "scheduler.bin", and "random_states.pkl"
            # If mixed precision was used, will also save a "scalar.bin" file
            if checkpointing_steps == "epoch":
                  A    =				f'''epoch_{epoch}'''
                  if args.output_dir is not None:
                        A    =				os.path.join(args.output_dir,    lowerCAmelCase     )
                  accelerator.save_state(lowerCAmelCase     )


def 			__UpperCamelCase    ()     ->  Any:
      A    =				argparse.ArgumentParser(description='Simple example of training script.'     )
      parser.add_argument(
          '--mixed_precision',    type=lowerCAmelCase,    default=lowerCAmelCase,    choices=['no', 'fp16', 'bf16', 'fp8'],    help='Whether to use mixed precision. Choose'
          'between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >= 1.10.'
          'and an Nvidia Ampere GPU.',    )
      parser.add_argument('--cpu',    action='store_true',    help='If passed, will train on the CPU.'     )
      parser.add_argument(
          '--checkpointing_steps',    type=lowerCAmelCase,    default=lowerCAmelCase,    help='Whether the various states should be saved at the end of every n steps, or \'epoch\' for each epoch.',    )
      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(
          '--resume_from_checkpoint',    type=lowerCAmelCase,    default=lowerCAmelCase,    help='If the training should continue from a checkpoint folder.',    )
      A    =				parser.parse_args()
      A    =				{'lr': 2E-5, 'num_epochs': 3, 'seed': 42, 'batch_size': 16}
      training_function(lowerCAmelCase,    lowerCAmelCase     )


if __name__ == "__main__":
  main()






import gc
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModelWithProjection, CLIPTokenizer

from diffusers import PriorTransformer, UnCLIPPipeline, UnCLIPScheduler, UNetaDConditionModel, UNetaDModel
from diffusers.pipelines.unclip.text_proj import UnCLIPTextProjModel
from diffusers.utils import load_numpy, nightly, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu, skip_mps

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineTesterMixin, assert_mean_pixel_difference


enable_full_determinism()



class        _UpperCAmelCase (   __lowercase     ,				unittest.TestCase			):
  '''simple docstring'''







  SCREAMING_SNAKE_CASE							:				List[str]					      =				UnCLIPPipeline
  SCREAMING_SNAKE_CASE							:				List[str]					      =				TEXT_TO_IMAGE_PARAMS - {
      '''negative_prompt''',
      '''height''',
      '''width''',
      '''negative_prompt_embeds''',
      '''guidance_scale''',
      '''prompt_embeds''',
      '''cross_attention_kwargs''',
  }
  SCREAMING_SNAKE_CASE							:				int					      =				TEXT_TO_IMAGE_BATCH_PARAMS
  SCREAMING_SNAKE_CASE							:				Optional[int]					      =				[
      '''generator''',
      '''return_dict''',
      '''prior_num_inference_steps''',
      '''decoder_num_inference_steps''',
      '''super_res_num_inference_steps''',
  ]
  SCREAMING_SNAKE_CASE							:				str					      =				False







  @property
  def 	UpperCamelCase       ( self	: List[Any] ):
        return 32







  @property
  def 	UpperCamelCase       ( self	: int ):
        return 32







  @property
  def 	UpperCamelCase       ( self	: Tuple ):
        return self.time_input_dim







  @property
  def 	UpperCamelCase       ( self	: Union[str, Any] ):
        return self.time_input_dim * 4







  @property
  def 	UpperCamelCase       ( self	: List[Any] ):
        return 100







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







  @property
  def 	UpperCamelCase       ( self	: List[Any] ):
        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-0_5		,						num_attention_heads=4		,						num_hidden_layers=5		,						pad_token_id=1		,						vocab_size=1000		,						)
        return CLIPTextModelWithProjection(UpperCamelCase__ )







  @property
  def 	UpperCamelCase       ( self	: List[str] ):
        torch.manual_seed(0 )

        A    =				{
            'num_attention_heads': 2,
            'attention_head_dim': 12,
            'embedding_dim': self.text_embedder_hidden_size,
            'num_layers': 1,
        }

        A    =				PriorTransformer(**UpperCamelCase__ )
        return model







  @property
  def 	UpperCamelCase       ( self	: Optional[int] ):
        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(**UpperCamelCase__ )
        return model







  @property
  def 	UpperCamelCase       ( self	: List[Any] ):
        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(**UpperCamelCase__ )
        return model







  @property
  def 	UpperCamelCase       ( self	: Any ):
        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 	UpperCamelCase       ( self	: str ):
        torch.manual_seed(0 )

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







  @property
  def 	UpperCamelCase       ( self	: Any ):
        # 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 	UpperCamelCase       ( self	: List[Any] ):
        A    =				self.dummy_prior
        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='fixed_small_log'		,						prediction_type='sample'		,						num_train_timesteps=1000		,						clip_sample_range=5.0		,						)

        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    =				{
            'prior': prior,
            'decoder': decoder,
            'text_proj': text_proj,
            'text_encoder': text_encoder,
            'tokenizer': tokenizer,
            'super_res_first': super_res_first,
            'super_res_last': super_res_last,
            'prior_scheduler': prior_scheduler,
            'decoder_scheduler': decoder_scheduler,
            'super_res_scheduler': super_res_scheduler,
        }

        return components







  def 	UpperCamelCase       ( self	: List[Any]		,						UpperCamelCase__	: Optional[Any]		,						UpperCamelCase__	: Optional[int]=0 ):
        if str(UpperCamelCase__ ).startswith('mps' ):
              A    =				torch.manual_seed(UpperCamelCase__ )
        else:
              A    =				torch.Generator(device=UpperCamelCase__ ).manual_seed(UpperCamelCase__ )
        A    =				{
            'prompt': 'horse',
            'generator': generator,
            'prior_num_inference_steps': 2,
            'decoder_num_inference_steps': 2,
            'super_res_num_inference_steps': 2,
            'output_type': 'numpy',
        }
        return inputs







  def 	UpperCamelCase       ( self	: int ):
        A    =				'cpu'

        A    =				self.get_dummy_components()

        A    =				self.pipeline_class(**UpperCamelCase__ )
        A    =				pipe.to(UpperCamelCase__ )

        pipe.set_progress_bar_config(disable=UpperCamelCase__ )

        A    =				pipe(**self.get_dummy_inputs(UpperCamelCase__ ) )
        A    =				output.images

        A    =				pipe(
            **self.get_dummy_inputs(UpperCamelCase__ )		,						return_dict=UpperCamelCase__		,						)[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.9_997,
                0.9_988,
                0.0_028,
                0.9_997,
                0.9_984,
                0.9_965,
                0.0_029,
                0.9_986,
                0.0_025,
            ] )

        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 	UpperCamelCase       ( self	: Any ):
        A    =				torch.device('cpu' )



        class        _UpperCAmelCase :
          '''simple docstring'''







          SCREAMING_SNAKE_CASE							:				Tuple					      =				1
        A    =				self.get_dummy_components()

        A    =				self.pipeline_class(**UpperCamelCase__ )
        A    =				pipe.to(UpperCamelCase__ )

        A    =				components['prior']
        A    =				components['decoder']
        A    =				components['super_res_first']
        A    =				components['tokenizer']
        A    =				components['text_encoder']

        A    =				torch.Generator(device=UpperCamelCase__ ).manual_seed(0 )
        A    =				prior.dtype
        A    =				1

        A    =				(batch_size, prior.config.embedding_dim)
        A    =				pipe.prepare_latents(
            UpperCamelCase__		,						dtype=UpperCamelCase__		,						device=UpperCamelCase__		,						generator=UpperCamelCase__		,						latents=UpperCamelCase__		,						scheduler=DummyScheduler() )
        A    =				(batch_size, decoder.config.in_channels, decoder.config.sample_size, decoder.config.sample_size)
        A    =				pipe.prepare_latents(
            UpperCamelCase__		,						dtype=UpperCamelCase__		,						device=UpperCamelCase__		,						generator=UpperCamelCase__		,						latents=UpperCamelCase__		,						scheduler=DummyScheduler() )

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

        pipe.set_progress_bar_config(disable=UpperCamelCase__ )

        A    =				'this is a prompt example'

        A    =				torch.Generator(device=UpperCamelCase__ ).manual_seed(0 )
        A    =				pipe(
            [prompt]		,						generator=UpperCamelCase__		,						prior_num_inference_steps=2		,						decoder_num_inference_steps=2		,						super_res_num_inference_steps=2		,						prior_latents=UpperCamelCase__		,						decoder_latents=UpperCamelCase__		,						super_res_latents=UpperCamelCase__		,						output_type='np'		,						)
        A    =				output.images

        A    =				tokenizer(
            UpperCamelCase__		,						padding='max_length'		,						max_length=tokenizer.model_max_length		,						return_tensors='pt'		,						)
        A    =				text_encoder(text_inputs.input_ids )
        A    =				text_inputs.attention_mask

        A    =				torch.Generator(device=UpperCamelCase__ ).manual_seed(0 )
        A    =				pipe(
            generator=UpperCamelCase__		,						prior_num_inference_steps=2		,						decoder_num_inference_steps=2		,						super_res_num_inference_steps=2		,						prior_latents=UpperCamelCase__		,						decoder_latents=UpperCamelCase__		,						super_res_latents=UpperCamelCase__		,						text_model_output=UpperCamelCase__		,						text_attention_mask=UpperCamelCase__		,						output_type='np'		,						)[0]

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







  @skip_mps
  def 	UpperCamelCase       ( self	: Any ):
        A    =				torch_device == 'cpu'

        self._test_attention_slicing_forward_pass(test_max_difference=UpperCamelCase__		,						expected_max_diff=0.01 )







  @skip_mps
  def 	UpperCamelCase       ( self	: Optional[Any] ):
        A    =				torch_device == 'cpu'
        A    =				True
        A    =				[
            'prior_num_inference_steps',
            'decoder_num_inference_steps',
            'super_res_num_inference_steps',
        ]

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







  def 	UpperCamelCase       ( self	: List[str] ):
        A    =				[
            'prior_num_inference_steps',
            '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=UpperCamelCase__		,						additional_params_copy_to_batched_inputs=UpperCamelCase__		,						)
        else:
              self._test_inference_batch_consistent(
                  additional_params_copy_to_batched_inputs=UpperCamelCase__ )







  @skip_mps
  def 	UpperCamelCase       ( self	: int ):
        return super().test_dict_tuple_outputs_equivalent()







  @skip_mps
  def 	UpperCamelCase       ( self	: Optional[Any] ):
        return super().test_save_load_local()







  @skip_mps
  def 	UpperCamelCase       ( self	: List[str] ):
        return super().test_save_load_optional_components()



@nightly
class        _UpperCAmelCase (   unittest.TestCase			):
  '''simple docstring'''







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







  def 	UpperCamelCase       ( self	: List[Any] ):
        A    =				load_numpy(
            'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
            '/unclip/karlo_v1_alpha_horse_cpu.npy' )

        A    =				UnCLIPPipeline.from_pretrained('kakaobrain/karlo-v1-alpha' )
        pipeline.set_progress_bar_config(disable=UpperCamelCase__ )

        A    =				torch.manual_seed(0 )
        A    =				pipeline(
            'horse'		,						num_images_per_prompt=1		,						generator=UpperCamelCase__		,						output_type='np'		,						)

        A    =				output.images[0]

        assert image.shape == (256, 256, 3)
        assert np.abs(expected_image - image ).max() < 1e-1



@slow
@require_torch_gpu
class        _UpperCAmelCase (   unittest.TestCase			):
  '''simple docstring'''







  def 	UpperCamelCase       ( self	: List[str] ):
        # clean up the VRAM after each test
        super().tearDown()
        gc.collect()
        torch.cuda.empty_cache()







  def 	UpperCamelCase       ( self	: List[Any] ):
        A    =				load_numpy(
            'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
            '/unclip/karlo_v1_alpha_horse_fp16.npy' )

        A    =				UnCLIPPipeline.from_pretrained('kakaobrain/karlo-v1-alpha'		,						torch_dtype=torch.floataa )
        A    =				pipeline.to(UpperCamelCase__ )
        pipeline.set_progress_bar_config(disable=UpperCamelCase__ )

        A    =				torch.Generator(device='cpu' ).manual_seed(0 )
        A    =				pipeline(
            'horse'		,						generator=UpperCamelCase__		,						output_type='np'		,						)

        A    =				output.images[0]

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

        assert_mean_pixel_difference(UpperCamelCase__		,						UpperCamelCase__ )







  def 	UpperCamelCase       ( self	: Dict ):
        torch.cuda.empty_cache()
        torch.cuda.reset_max_memory_allocated()
        torch.cuda.reset_peak_memory_stats()

        A    =				UnCLIPPipeline.from_pretrained('kakaobrain/karlo-v1-alpha'		,						torch_dtype=torch.floataa )
        A    =				pipe.to(UpperCamelCase__ )
        pipe.set_progress_bar_config(disable=UpperCamelCase__ )
        pipe.enable_attention_slicing()
        pipe.enable_sequential_cpu_offload()

        A    =				pipe(
            'horse'		,						num_images_per_prompt=1		,						prior_num_inference_steps=2		,						decoder_num_inference_steps=2		,						super_res_num_inference_steps=2		,						output_type='np'		,						)

        A    =				torch.cuda.max_memory_allocated()
        # make sure that less than 7 GB is allocated
        assert mem_bytes < 7 * 10**9



