

import timeit
from typing import Callable, Optional

from ..configuration_utils import PretrainedConfig
from ..models.auto.modeling_auto import MODEL_MAPPING, MODEL_WITH_LM_HEAD_MAPPING
from ..utils import is_pyanvml_available, is_torch_available, logging
from .benchmark_utils import (
    Benchmark,
    Memory,
    MemorySummary,
    measure_peak_memory_cpu,
    start_memory_tracing,
    stop_memory_tracing,
)


if is_torch_available():
  import torch

  from .benchmark_args import PyTorchBenchmarkArguments


if is_pyanvml_available():
  import pyanvml.pyanvml as nvml


_UpperCAmelCase    	=			logging.get_logger(__name__)



class        _UpperCAmelCase (   __lowercase			):
  '''simple docstring'''







  SCREAMING_SNAKE_CASE							:				PyTorchBenchmarkArguments
  SCREAMING_SNAKE_CASE							:				PretrainedConfig
  SCREAMING_SNAKE_CASE							:				str					      =				"PyTorch"







  @property
  def 	UpperCamelCase       ( self	: Optional[Any] ):
        return torch.__version__







  def 	UpperCamelCase       ( self	: Tuple		,						UpperCamelCase__	: str		,						UpperCamelCase__	: int		,						UpperCamelCase__	: int ):
        A    =				self._prepare_inference_func(UpperCamelCase__		,						UpperCamelCase__		,						UpperCamelCase__ )
        return self._measure_speed(_inference )







  def 	UpperCamelCase       ( self	: Dict		,						UpperCamelCase__	: str		,						UpperCamelCase__	: int		,						UpperCamelCase__	: int ):
        A    =				self._prepare_inference_func(UpperCamelCase__		,						UpperCamelCase__		,						UpperCamelCase__ )
        return self._measure_memory(_inference )







  def 	UpperCamelCase       ( self	: Optional[int]		,						UpperCamelCase__	: str		,						UpperCamelCase__	: int		,						UpperCamelCase__	: int ):
        A    =				self._prepare_train_func(UpperCamelCase__		,						UpperCamelCase__		,						UpperCamelCase__ )
        return self._measure_speed(_train )







  def 	UpperCamelCase       ( self	: Optional[Any]		,						UpperCamelCase__	: str		,						UpperCamelCase__	: int		,						UpperCamelCase__	: int ):
        A    =				self._prepare_train_func(UpperCamelCase__		,						UpperCamelCase__		,						UpperCamelCase__ )
        return self._measure_memory(_train )







  def 	UpperCamelCase       ( self	: Tuple		,						UpperCamelCase__	: str		,						UpperCamelCase__	: int		,						UpperCamelCase__	: int ):
        A    =				self.config_dict[model_name]

        if self.args.torchscript:
              A    =				True

        A    =				(
            hasattr(UpperCamelCase__		,						'architectures' )
            and isinstance(config.architectures		,						UpperCamelCase__ )
            and len(config.architectures ) > 0
        )
        if not self.args.only_pretrain_model and has_model_class_in_config:
              try:
                    A    =				config.architectures[0]
                    A    =				__import__('transformers'		,						fromlist=[model_class] )
                    A    =				getattr(UpperCamelCase__		,						UpperCamelCase__ )
                    A    =				model_cls(UpperCamelCase__ )
              except ImportError:
                    raise ImportError(
                        f'''{model_class} does not exist. If you just want to test the pretrained model, you might want to'''
                        ' set `--only_pretrain_model` or `args.only_pretrain_model=True`.' )
        else:
              A    =				MODEL_MAPPING[config.__class__](UpperCamelCase__ )

        model.eval()
        model.to(self.args.device )

        # encoder-decoder has vocab size saved differently
        A    =				config.vocab_size if hasattr(UpperCamelCase__		,						'vocab_size' ) else config.encoder.vocab_size
        A    =				torch.randint(UpperCamelCase__		,						(batch_size, sequence_length)		,						dtype=torch.long		,						device=self.args.device )

        if self.args.fpaa:
              logger.info('Running training in Mixed Precision...' )
              if not self.args.is_gpu:
                    raise ValueError('Mixed precision is possible only for GPU.' )
              # amp seems to have memory leaks so that memory usage
              # is measured using .half() for now https://github.com/NVIDIA/apex/issues/439
              model.half()

        if self.args.torchscript:
              with torch.no_grad():
                    A    =				torch.jit.trace(UpperCamelCase__		,						UpperCamelCase__ )
        else:
              A    =				model

        def encoder_decoder_forward():
              with torch.no_grad():
                    A    =				inference_model(UpperCamelCase__		,						decoder_input_ids=UpperCamelCase__ )
              return outputs

        def encoder_forward():
              with torch.no_grad():
                    A    =				inference_model(UpperCamelCase__ )
              return outputs

        A    =				encoder_decoder_forward if config.is_encoder_decoder else encoder_forward
        return _forward







  def 	UpperCamelCase       ( self	: Optional[Any]		,						UpperCamelCase__	: str		,						UpperCamelCase__	: int		,						UpperCamelCase__	: int ):
        A    =				self.config_dict[model_name]

        A    =				(
            hasattr(UpperCamelCase__		,						'architectures' )
            and isinstance(config.architectures		,						UpperCamelCase__ )
            and len(config.architectures ) > 0
        )
        if not self.args.only_pretrain_model and has_model_class_in_config:
              try:
                    A    =				config.architectures[0]
                    A    =				__import__('transformers'		,						fromlist=[model_class] )
                    A    =				getattr(UpperCamelCase__		,						UpperCamelCase__ )
                    A    =				model_cls(UpperCamelCase__ )
              except ImportError:
                    raise ImportError(
                        f'''{model_class} does not exist. If you just want to test the pretrained model, you might want to'''
                        ' set `--only_pretrain_model` or `args.only_pretrain_model=True`.' )
        else:
              A    =				MODEL_WITH_LM_HEAD_MAPPING[config.__class__](UpperCamelCase__ )

        if self.args.torchscript:
              raise NotImplementedError('Training for torchscript is currently not implemented' )
        else:
              A    =				model

        model.train()
        model.to(self.args.device )

        # encoder-decoder has vocab size saved differently
        A    =				config.vocab_size if hasattr(UpperCamelCase__		,						'vocab_size' ) else config.encoder.vocab_size
        A    =				torch.randint(UpperCamelCase__		,						(batch_size, sequence_length)		,						dtype=torch.long		,						device=self.args.device )

        if self.args.fpaa:
              logger.info('Running training in Mixed Precision...' )
              if not self.args.is_gpu:
                    raise ValueError('Mixed precision is possible only for GPU.' )

              # amp seems to have memory leaks so that memory usage
              # is measured using .half() for now https://github.com/NVIDIA/apex/issues/439
              model.half()

        def compute_loss_and_backprob_encoder():
              A    =				train_model(UpperCamelCase__		,						labels=UpperCamelCase__ )[0]
              loss.backward()
              return loss

        def compute_loss_and_backprob_encoder_decoder():
              A    =				train_model(UpperCamelCase__		,						decoder_input_ids=UpperCamelCase__		,						labels=UpperCamelCase__ )[0]
              loss.backward()
              return loss

        A    =				(
            compute_loss_and_backprob_encoder_decoder
            if config.is_encoder_decoder
            else compute_loss_and_backprob_encoder
        )
        return _train







  def 	UpperCamelCase       ( self	: List[Any]		,						UpperCamelCase__	: str ):
        try:
              if self.args.is_tpu or self.args.torchscript:
                    # run additional 10 times to stabilize compilation for tpu and torchscript
                    logger.info('Do inference on TPU or torchscript. Running model 5 times to stabilize compilation' )
                    timeit.repeat(
                        UpperCamelCase__		,						repeat=1		,						number=5		,						)

              # as written in https://docs.python.org/2/library/timeit.html#timeit.Timer.repeat, min should be taken rather than the average
              A    =				timeit.repeat(
                  UpperCamelCase__		,						repeat=self.args.repeat		,						number=10		,						)

              if self.args.is_tpu and self.args.torch_xla_tpu_print_metrics:
                    import torch_xla.debug.metrics as met

                    self.print_fn(met.metrics_report() )

              return min(UpperCamelCase__ ) / 10.0
        except RuntimeError as e:
              self.print_fn(f'''Doesn\'t fit on GPU. {e}''' )
              return "N/A"







  def 	UpperCamelCase       ( self	: Tuple		,						UpperCamelCase__	: Callable[[], None] ):
        try:
              if self.args.trace_memory_line_by_line:
                    A    =				start_memory_tracing('transformers' )

              if self.args.is_tpu:
                    # tpu
                    raise NotImplementedError(
                        'Memory Benchmarking is currently not implemented for TPU. Please disable memory benchmarking with'
                        ' `--no-memory` or `args.memory=False`' )
              elif self.args.is_gpu:
                    if not is_pyanvml_available():
                          logger.warning(
                              'py3nvml not installed, we won\'t log GPU memory usage. '
                              'Install py3nvml (pip install py3nvml) to log information about GPU.' )
                          A    =				'N/A'
                    else:
                          logger.info(
                              'Measuring total GPU usage on GPU device. Make sure to not have additional processes running'
                              ' on the same GPU.' )
                          # init nvml
                          nvml.nvmlInit()
                          func()
                          A    =				nvml.nvmlDeviceGetHandleByIndex(self.args.device_idx )
                          A    =				nvml.nvmlDeviceGetMemoryInfo(UpperCamelCase__ )
                          A    =				meminfo.used
                          A    =				Memory(UpperCamelCase__ )
                          # shutdown nvml
                          nvml.nvmlShutdown()
              else:
                    # cpu
                    A    =				measure_peak_memory_cpu(UpperCamelCase__ )
                    A    =				Memory(UpperCamelCase__ ) if isinstance(UpperCamelCase__		,						UpperCamelCase__ ) else memory_bytes

              if self.args.trace_memory_line_by_line:
                    A    =				stop_memory_tracing(UpperCamelCase__ )
              else:
                    A    =				None

              return memory, summary
        except RuntimeError as e:
              self.print_fn(f'''Doesn\'t fit on GPU. {e}''' )
              return "N/A", None






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



