

import argparse
import json
import os.path
from collections import OrderedDict

import numpy as np
import requests
import torch
from flax.training.checkpoints import restore_checkpoint
from huggingface_hub import hf_hub_download

from transformers import VivitConfig, VivitForVideoClassification, VivitImageProcessor
from transformers.image_utils import PILImageResampling


def 			__UpperCamelCase    (lowerCAmelCase					:  int     )     ->  Optional[Any]:
      A    =				'https://storage.googleapis.com/scenic-bucket/vivit/kinetics_400/vivit_base_16x2_unfactorized/checkpoint'

      with open(lowerCAmelCase,    'wb'     ) as f:
            with requests.get(lowerCAmelCase,    stream=lowerCAmelCase     ) as req:
                  for chunk in req.iter_content(chunk_size=2_048     ):
                        f.write(lowerCAmelCase     )


def 			__UpperCamelCase    ()     ->  VivitConfig:
      A    =				VivitConfig()

      A    =				400
      A    =				'huggingface/label-files'
      A    =				'kinetics400-id2label.json'

      A    =				json.load(open(hf_hub_download(lowerCAmelCase,    lowerCAmelCase,    repo_type='dataset'     ),    'r'     )     )
      A    =				{int(lowerCAmelCase     ): v for k, v in idalabel.items()}
      A    =				idalabel
      A    =				{v: k for k, v in idalabel.items()}
      return config


def 			__UpperCamelCase    ()     ->  int:
      A    =				hf_hub_download(
          repo_id='hf-internal-testing/spaghetti-video',    filename='eating_spaghetti_32_frames.npy',    repo_type='dataset'     )
      A    =				np.load(lowerCAmelCase     )
      return list(lowerCAmelCase     )


def 			__UpperCamelCase    (lowerCAmelCase					:  np.ndarray     )     ->  Union[str, Any]:
      if np.ndim(lowerCAmelCase     ) == 2:
            return transform_attention_bias(lowerCAmelCase     )

      elif np.ndim(lowerCAmelCase     ) == 3:
            return transform_attention_kernel(lowerCAmelCase     )

      else:
            raise Exception(f'''Invalid number of dimesions: {np.ndim(lowerCAmelCase     )}'''     )


def 			__UpperCamelCase    (lowerCAmelCase					:  np.ndarray     )     ->  Optional[Any]:
      return current.flatten()


def 			__UpperCamelCase    (lowerCAmelCase					:  np.ndarray     )     ->  List[str]:
      return np.reshape(lowerCAmelCase,    (current.shape[0], current.shape[1] * current.shape[2])     ).T


def 			__UpperCamelCase    (lowerCAmelCase					:  np.ndarray     )     ->  Dict:
      return np.reshape(lowerCAmelCase,    (current.shape[0] * current.shape[1], current.shape[2])     ).T


def 			__UpperCamelCase    (lowerCAmelCase					:  int,    lowerCAmelCase					:  Any     )     ->  int:
      A    =				state_dict['optimizer']['target']['Transformer'][f'''encoderblock_{i}''']

      A    =				f'''encoder.layer.{i}.'''
      A    =				{
          prefix + 'intermediate.dense.bias': state['MlpBlock_0']['Dense_0']['bias'],
          prefix + 'intermediate.dense.weight': np.transpose(state['MlpBlock_0']['Dense_0']['kernel']     ),
          prefix + 'output.dense.bias': state['MlpBlock_0']['Dense_1']['bias'],
          prefix + 'output.dense.weight': np.transpose(state['MlpBlock_0']['Dense_1']['kernel']     ),
          prefix + 'layernorm_before.bias': state['LayerNorm_0']['bias'],
          prefix + 'layernorm_before.weight': state['LayerNorm_0']['scale'],
          prefix + 'layernorm_after.bias': state['LayerNorm_1']['bias'],
          prefix + 'layernorm_after.weight': state['LayerNorm_1']['scale'],
          prefix
          + 'attention.attention.query.bias': transform_attention(
              state['MultiHeadDotProductAttention_0']['query']['bias']     ),
          prefix
          + 'attention.attention.query.weight': transform_attention(
              state['MultiHeadDotProductAttention_0']['query']['kernel']     ),
          prefix
          + 'attention.attention.key.bias': transform_attention(state['MultiHeadDotProductAttention_0']['key']['bias']     ),
          prefix
          + 'attention.attention.key.weight': transform_attention(
              state['MultiHeadDotProductAttention_0']['key']['kernel']     ),
          prefix
          + 'attention.attention.value.bias': transform_attention(
              state['MultiHeadDotProductAttention_0']['value']['bias']     ),
          prefix
          + 'attention.attention.value.weight': transform_attention(
              state['MultiHeadDotProductAttention_0']['value']['kernel']     ),
          prefix + 'attention.output.dense.bias': state['MultiHeadDotProductAttention_0']['out']['bias'],
          prefix
          + 'attention.output.dense.weight': transform_attention_output_weight(
              state['MultiHeadDotProductAttention_0']['out']['kernel']     ),
      }

      return new_state


def 			__UpperCamelCase    (lowerCAmelCase					:  Tuple     )     ->  Optional[int]:
      return sum([1 if 'encoderblock_' in k else 0 for k in state_dict['optimizer']['target']['Transformer'].keys()]     )


def 			__UpperCamelCase    (lowerCAmelCase					:  Optional[Any],    lowerCAmelCase					:  Optional[int]=False     )     ->  Optional[int]:
      A    =				get_n_layers(lowerCAmelCase     )

      A    =				OrderedDict()

      A    =				state_dict['optimizer']['target']['Transformer']['encoder_norm']['bias']
      A    =				state_dict['optimizer']['target']['Transformer']['encoder_norm']['scale']

      A    =				np.transpose(
          state_dict['optimizer']['target']['embedding']['kernel'],    (4, 3, 0, 1, 2)     )
      A    =				state_dict['optimizer']['target']['embedding']['bias']

      A    =				state_dict['optimizer']['target']['cls']
      A    =				state_dict['optimizer']['target']['Transformer']['posembed_input'][
          'pos_embedding'
      ]

      for i in range(lowerCAmelCase     ):
            new_state.update(transform_state_encoder_block(lowerCAmelCase,    lowerCAmelCase     )     )

      if classification_head:
            A    =				{'vivit.' + k: v for k, v in new_state.items()}
            A    =				np.transpose(state_dict['optimizer']['target']['output_projection']['kernel']     )
            A    =				np.transpose(state_dict['optimizer']['target']['output_projection']['bias']     )

      return {k: torch.tensor(lowerCAmelCase     ) for k, v in new_state.items()}


def 			__UpperCamelCase    ()     ->  VivitImageProcessor:
      A    =				VivitImageProcessor()

      assert extractor.do_resize is True
      assert extractor.size == {"shortest_edge": 256}
      assert extractor.do_center_crop is True
      assert extractor.crop_size == {"width": 224, "height": 224}
      assert extractor.resample == PILImageResampling.BILINEAR

      # here: https://github.com/deepmind/dmvr/blob/master/dmvr/modalities.py
      # one can seen that add_image has default values for normalization_mean and normalization_std set to 0 and 1
      # which effectively means no normalization (and ViViT does not overwrite those when calling this func)
      assert extractor.do_normalize is False
      assert extractor.do_rescale is True
      assert extractor.rescale_factor == 1 / 255

      # zero-centering = True in original implementation
      assert extractor.do_zero_centering is True

      return extractor


def 			__UpperCamelCase    (lowerCAmelCase					:  str     )     ->  Optional[int]:
      A    =				'checkpoint'

      if not os.path.exists(lowerCAmelCase     ):
            download_checkpoint(lowerCAmelCase     )

      A    =				restore_checkpoint(lowerCAmelCase,    lowerCAmelCase     )
      A    =				transform_state(lowerCAmelCase,    classification_head=lowerCAmelCase     )

      A    =				get_vivit_config()

      assert config.image_size == 224
      assert config.num_frames == 32

      A    =				VivitForVideoClassification(lowerCAmelCase     )
      model.load_state_dict(lowerCAmelCase     )
      model.eval()

      A    =				get_processor()

      A    =				prepare_video()
      A    =				extractor(lowerCAmelCase,    return_tensors='pt'     )

      A    =				model(**lowerCAmelCase     )

      A    =				torch.Size([1, 400]     )
      A    =				torch.tensor([-1.0543, 2.0764, -0.2104, 0.4439, -0.9658]     )

      assert outputs.logits.shape == expected_shape
      assert torch.allclose(outputs.logits[0, :5],    lowerCAmelCase,    atol=1E-4     ), outputs.logits[0, :5]

      model.save_pretrained(lowerCAmelCase     )
      extractor.save_pretrained(lowerCAmelCase     )


if __name__ == "__main__":
  _UpperCAmelCase    	=			argparse.ArgumentParser()

  parser.add_argument("--output_model_name", "-o", type=str, help="Output path for the converted HuggingFace model")

  _UpperCAmelCase    	=			parser.parse_args()
  convert(args.output_model_name)






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



