import collections
import gzip
import os
import urllib

import numpy
from tensorflow.python.framework import dtypes, random_seed
from tensorflow.python.platform import gfile
from tensorflow.python.util.deprecation import deprecated

UpperCAmelCase_			    =							collections.namedtuple('''_Datasets''', ['''train''', '''validation''', '''test'''])

# CVDF mirror of http://yann.lecun.com/exdb/mnist/
UpperCAmelCase_			    =							'''https://storage.googleapis.com/cvdf-datasets/mnist/'''




def       lowerCAmelCase_     (		lowercase:							int		)    ->				Optional[Any]:







 '''simple docstring'''
 _UpperCamelCase:     Dict							= numpy.dtype(numpy.uintaa		).newbyteorder('''>'''		)
 return numpy.frombuffer(bytestream.read(4		)  ,					dtype=lowercase		)[0]




@deprecated(lowercase  ,					'''Please use tf.data to implement this functionality.'''		)
def       lowerCAmelCase_     (		lowercase:							Optional[Any]		)    ->				Optional[int]:







 '''simple docstring'''
 print('''Extracting'''  ,					f.name		)
 with gzip.GzipFile(fileobj=lowercase		) as bytestream:
  _UpperCamelCase:     List[str]							= _readaa(lowercase		)
  if magic != 2_051:
   raise ValueError(
       '''Invalid magic number %d in MNIST image file: %s''' % (magic, f.name)		)
  _UpperCamelCase:     str							= _readaa(lowercase		)
  _UpperCamelCase:     Dict							= _readaa(lowercase		)
  _UpperCamelCase:     List[Any]							= _readaa(lowercase		)
  _UpperCamelCase:     str							= bytestream.read(rows * cols * num_images		)
  _UpperCamelCase:     List[str]							= numpy.frombuffer(lowercase  ,					dtype=numpy.uinta		)
  _UpperCamelCase:     List[Any]							= data.reshape(lowercase  ,					lowercase  ,					lowercase  ,					1		)
  return data




@deprecated(lowercase  ,					'''Please use tf.one_hot on tensors.'''		)
def       lowerCAmelCase_     (		lowercase:							List[Any]  ,					lowercase:							int		)    ->				Optional[int]:







 '''simple docstring'''
 _UpperCamelCase:     str							= labels_dense.shape[0]
 _UpperCamelCase:     Union[str, Any]							= numpy.arange(lowercase		) * num_classes
 _UpperCamelCase:     Optional[Any]							= numpy.zeros((num_labels, num_classes)		)
 _UpperCamelCase:     Union[str, Any]							= 1
 return labels_one_hot




@deprecated(lowercase  ,					'''Please use tf.data to implement this functionality.'''		)
def       lowerCAmelCase_     (		lowercase:							Dict  ,					lowercase:							List[Any]=False  ,					lowercase:							Dict=10		)    ->				Optional[Any]:







 '''simple docstring'''
 print('''Extracting'''  ,					f.name		)
 with gzip.GzipFile(fileobj=lowercase		) as bytestream:
  _UpperCamelCase:     Dict							= _readaa(lowercase		)
  if magic != 2_049:
   raise ValueError(
       '''Invalid magic number %d in MNIST label file: %s''' % (magic, f.name)		)
  _UpperCamelCase:     Union[str, Any]							= _readaa(lowercase		)
  _UpperCamelCase:     Optional[Any]							= bytestream.read(lowercase		)
  _UpperCamelCase:     Optional[int]							= numpy.frombuffer(lowercase  ,					dtype=numpy.uinta		)
  if one_hot:
   return _dense_to_one_hot(lowercase  ,					lowercase		)
  return labels




class  __magic_name__  :





   """simple docstring"""




   @deprecated(
       _lowercase		,  '''Please use alternatives such as official/mnist/_DataSet.py'''
       ''' from tensorflow/models.'''		,  )
   def __init__(	self					:  Dict		,  _lowercase					:  List[str]		,  _lowercase					:  Optional[int]		,  _lowercase					:  Optional[Any]=False		,  _lowercase					:  str=False		,  _lowercase					:  Optional[int]=dtypes.floataa		,  _lowercase					:  Optional[int]=True		,  _lowercase					:  List[str]=None		,  ):



    """simple docstring"""



    _UpperCamelCase	,							_UpperCamelCase:     Tuple							= random_seed.get_seed(_lowercase							)
    # If op level seed is not set, use whatever graph level seed is returned
    numpy.random.seed(seeda if seed is None else seeda							)
    _UpperCamelCase:     Optional[Any]							= dtypes.as_dtype(_lowercase							).base_dtype
    if dtype not in (dtypes.uinta, dtypes.floataa):
     raise TypeError('''Invalid image dtype %r, expected uint8 or float32''' % dtype							)
    if fake_data:
     _UpperCamelCase:     List[Any]							= 10_000
     _UpperCamelCase:     List[str]							= one_hot
    else:
     assert (
         images.shape[0] == labels.shape[0]
     ), f"""images.shape: {images.shape} labels.shape: {labels.shape}"""
     _UpperCamelCase:     List[str]							= images.shape[0]

     # Convert shape from [num examples, rows, columns, depth]
     # to [num examples, rows*columns] (assuming depth == 1)
     if reshape:
      assert images.shape[3] == 1
      _UpperCamelCase:     Optional[Any]							= images.reshape(
          images.shape[0]		,  images.shape[1] * images.shape[2]							)
     if dtype == dtypes.floataa:
      # Convert from [0, 255] -> [0.0, 1.0].
      _UpperCamelCase:     Optional[Any]							= images.astype(numpy.floataa							)
      _UpperCamelCase:     str							= numpy.multiply(_lowercase		,  1.0 / 255.0							)
    _UpperCamelCase:     Dict							= images
    _UpperCamelCase:     List[Any]							= labels
    _UpperCamelCase:     int							= 0
    _UpperCamelCase:     str							= 0



   @property
   def   lowerCAmelCase						(	self					:  int							):



    """simple docstring"""



    return self._images



   @property
   def   lowerCAmelCase						(	self					:  Optional[Any]							):



    """simple docstring"""



    return self._labels



   @property
   def   lowerCAmelCase						(	self					:  Any							):



    """simple docstring"""



    return self._num_examples



   @property
   def   lowerCAmelCase						(	self					:  Tuple							):



    """simple docstring"""



    return self._epochs_completed



   def   lowerCAmelCase						(	self					:  int		,  _lowercase					:  List[str]		,  _lowercase					:  int=False		,  _lowercase					:  Optional[int]=True							):



    """simple docstring"""



    if fake_data:
     _UpperCamelCase:     int							= [1] * 784
     _UpperCamelCase:     int							= [1] + [0] * 9 if self.one_hot else 0
     return (
         [fake_image for _ in range(_lowercase							)],
         [fake_label for _ in range(_lowercase							)],
     )
    _UpperCamelCase:     int							= self._index_in_epoch
    # Shuffle for the first epoch
    if self._epochs_completed == 0 and start == 0 and shuffle:
     _UpperCamelCase:     Optional[Any]							= numpy.arange(self._num_examples							)
     numpy.random.shuffle(_lowercase							)
     _UpperCamelCase:     Dict							= self.images[perma]
     _UpperCamelCase:     Tuple							= self.labels[perma]
    # Go to the next epoch
    if start + batch_size > self._num_examples:
     # Finished epoch
     self._epochs_completed += 1
     # Get the rest examples in this epoch
     _UpperCamelCase:     Any							= self._num_examples - start
     _UpperCamelCase:     str							= self._images[start : self._num_examples]
     _UpperCamelCase:     Dict							= self._labels[start : self._num_examples]
     # Shuffle the data
     if shuffle:
      _UpperCamelCase:     Optional[Any]							= numpy.arange(self._num_examples							)
      numpy.random.shuffle(_lowercase							)
      _UpperCamelCase:     Tuple							= self.images[perm]
      _UpperCamelCase:     Tuple							= self.labels[perm]
     # Start next epoch
     _UpperCamelCase:     Optional[int]							= 0
     _UpperCamelCase:     Optional[Any]							= batch_size - rest_num_examples
     _UpperCamelCase:     Any							= self._index_in_epoch
     _UpperCamelCase:     Optional[Any]							= self._images[start:end]
     _UpperCamelCase:     Tuple							= self._labels[start:end]
     return (
         numpy.concatenate((images_rest_part, images_new_part)		,  axis=0							),
         numpy.concatenate((labels_rest_part, labels_new_part)		,  axis=0							),
     )
    else:
     self._index_in_epoch += batch_size
     _UpperCamelCase:     Dict							= self._index_in_epoch
     return self._images[start:end], self._labels[start:end]




@deprecated(lowercase  ,					'''Please write your own downloading logic.'''		)
def       lowerCAmelCase_     (		lowercase:							Tuple  ,					lowercase:							Dict  ,					lowercase:							Any		)    ->				int:







 '''simple docstring'''
 if not gfile.Exists(lowercase		):
  gfile.MakeDirs(lowercase		)
 _UpperCamelCase:     Any							= os.path.join(lowercase  ,					lowercase		)
 if not gfile.Exists(lowercase		):
  urllib.request.urlretrieve(lowercase  ,					lowercase		)  # noqa: S310
  with gfile.GFile(lowercase		) as f:
   _UpperCamelCase:     List[str]							= f.size()
  print('''Successfully downloaded'''  ,					lowercase  ,					lowercase  ,					'''bytes.'''		)
 return filepath





@deprecated(
    lowercase  ,					'''Please use alternatives such as:''' ''' tensorflow_datasets.load(\'mnist\')'''		)
def       lowerCAmelCase_     (		lowercase:							Any  ,					lowercase:							int=False  ,					lowercase:							Dict=False  ,					lowercase:							int=dtypes.floataa  ,					lowercase:							List[Any]=True  ,					lowercase:							Tuple=5_000  ,					lowercase:							Optional[int]=None  ,					lowercase:							Tuple=DEFAULT_SOURCE_URL  ,					)    ->				int:







 '''simple docstring'''
 if fake_data:

  def fake():
   return _DataSet(
       []  ,					[]  ,					fake_data=lowercase  ,					one_hot=lowercase  ,					dtype=lowercase  ,					seed=lowercase		)

  _UpperCamelCase:     Any							= fake()
  _UpperCamelCase:     Dict							= fake()
  _UpperCamelCase:     Optional[int]							= fake()
  return _Datasets(train=lowercase  ,					validation=lowercase  ,					test=lowercase		)

 if not source_url:  # empty string check
  _UpperCamelCase:     str							= DEFAULT_SOURCE_URL

 _UpperCamelCase:     Optional[int]							= '''train-images-idx3-ubyte.gz'''
 _UpperCamelCase:     int							= '''train-labels-idx1-ubyte.gz'''
 _UpperCamelCase:     Optional[Any]							= '''t10k-images-idx3-ubyte.gz'''
 _UpperCamelCase:     List[Any]							= '''t10k-labels-idx1-ubyte.gz'''

 _UpperCamelCase:     str							= _maybe_download(
     lowercase  ,					lowercase  ,					source_url + train_images_file		)
 with gfile.Open(lowercase  ,					'''rb'''		) as f:
  _UpperCamelCase:     List[str]							= _extract_images(lowercase		)

 _UpperCamelCase:     Optional[int]							= _maybe_download(
     lowercase  ,					lowercase  ,					source_url + train_labels_file		)
 with gfile.Open(lowercase  ,					'''rb'''		) as f:
  _UpperCamelCase:     Union[str, Any]							= _extract_labels(lowercase  ,					one_hot=lowercase		)

 _UpperCamelCase:     Tuple							= _maybe_download(
     lowercase  ,					lowercase  ,					source_url + test_images_file		)
 with gfile.Open(lowercase  ,					'''rb'''		) as f:
  _UpperCamelCase:     Union[str, Any]							= _extract_images(lowercase		)

 _UpperCamelCase:     Dict							= _maybe_download(
     lowercase  ,					lowercase  ,					source_url + test_labels_file		)
 with gfile.Open(lowercase  ,					'''rb'''		) as f:
  _UpperCamelCase:     Optional[Any]							= _extract_labels(lowercase  ,					one_hot=lowercase		)

 if not 0 <= validation_size <= len(lowercase		):
  _UpperCamelCase:     Dict							= (
      '''Validation size should be between 0 and '''
      F"""{len(lowercase		)}. Received: {validation_size}."""
  )
  raise ValueError(lowercase		)

 _UpperCamelCase:     Optional[int]							= train_images[:validation_size]
 _UpperCamelCase:     Tuple							= train_labels[:validation_size]
 _UpperCamelCase:     Optional[int]							= train_images[validation_size:]
 _UpperCamelCase:     Union[str, Any]							= train_labels[validation_size:]

 _UpperCamelCase:     Dict							= {'''dtype''': dtype, '''reshape''': reshape, '''seed''': seed}

 _UpperCamelCase:     Optional[Any]							= _DataSet(lowercase  ,					lowercase  ,					**lowercase		)
 _UpperCamelCase:     Dict							= _DataSet(lowercase  ,					lowercase  ,					**lowercase		)
 _UpperCamelCase:     str							= _DataSet(lowercase  ,					lowercase  ,					**lowercase		)

 return _Datasets(train=lowercase  ,					validation=lowercase  ,					test=lowercase		)
import gc
import random
import unittest

import numpy as np
import torch
from PIL import Image
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import AutoencoderKL, PNDMScheduler, StableDiffusionInpaintPipeline, UNetaDConditionModel
from diffusers.utils import floats_tensor, load_image, load_numpy, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu, slow

from ..pipeline_params import TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS, TEXT_GUIDED_IMAGE_INPAINTING_PARAMS
from ..test_pipelines_common import PipelineKarrasSchedulerTesterMixin, PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()




class  __magic_name__  (		__a      ,	__a      ,	__a      ,	unittest.TestCase    ):





   """simple docstring"""




   lowerCAmelCase      :    Optional[int]  			=   StableDiffusionInpaintPipeline
   lowerCAmelCase      :    Optional[Any]  			=   TEXT_GUIDED_IMAGE_INPAINTING_PARAMS
   lowerCAmelCase      :    str  			=   TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS
   lowerCAmelCase      :    str  			=   frozenset(
       []    )  # TO-DO: update image_params once pipeline is refactored with VaeImageProcessor.preprocess
   lowerCAmelCase      :    Tuple  			=   frozenset([]    )



   def   lowerCAmelCase						(	self					:  Dict							):



    """simple docstring"""



    torch.manual_seed(0							)
    _UpperCamelCase:     List[str]							= UNetaDConditionModel(
        block_out_channels=(32, 64)		,  layers_per_block=2		,  sample_size=32		,  in_channels=9		,  out_channels=4		,  down_block_types=('''DownBlock2D''', '''CrossAttnDownBlock2D''')		,  up_block_types=('''CrossAttnUpBlock2D''', '''UpBlock2D''')		,  cross_attention_dim=32		,  attention_head_dim=(2, 4)		,  use_linear_projection=_lowercase		,  )
    _UpperCamelCase:     List[str]							= PNDMScheduler(skip_prk_steps=_lowercase							)
    torch.manual_seed(0							)
    _UpperCamelCase:     List[str]							= AutoencoderKL(
        block_out_channels=[32, 64]		,  in_channels=3		,  out_channels=3		,  down_block_types=['''DownEncoderBlock2D''', '''DownEncoderBlock2D''']		,  up_block_types=['''UpDecoderBlock2D''', '''UpDecoderBlock2D''']		,  latent_channels=4		,  sample_size=128		,  )
    torch.manual_seed(0							)
    _UpperCamelCase:     Optional[int]							= CLIPTextConfig(
        bos_token_id=0		,  eos_token_id=2		,  hidden_size=32		,  intermediate_size=37		,  layer_norm_eps=1E-05		,  num_attention_heads=4		,  num_hidden_layers=5		,  pad_token_id=1		,  vocab_size=1_000		,  hidden_act='''gelu'''		,  projection_dim=512		,  )
    _UpperCamelCase:     Optional[int]							= CLIPTextModel(_lowercase							)
    _UpperCamelCase:     Union[str, Any]							= CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''							)

    _UpperCamelCase:     Optional[Any]							= {
        '''unet''': unet,
        '''scheduler''': scheduler,
        '''vae''': vae,
        '''text_encoder''': text_encoder,
        '''tokenizer''': tokenizer,
        '''safety_checker''': None,
        '''feature_extractor''': None,
    }
    return components



   def   lowerCAmelCase						(	self					:  Optional[int]		,  _lowercase					:  Any		,  _lowercase					:  int=0							):



    """simple docstring"""



    _UpperCamelCase:     Any							= floats_tensor((1, 3, 32, 32)		,  rng=random.Random(_lowercase							)							).to(_lowercase							)
    _UpperCamelCase:     Union[str, Any]							= image.cpu().permute(0		,  2		,  3		,  1							)[0]
    _UpperCamelCase:     Optional[int]							= Image.fromarray(np.uinta(_lowercase							)							).convert('''RGB'''							).resize((64, 64)							)
    _UpperCamelCase:     Optional[Any]							= Image.fromarray(np.uinta(image + 4							)							).convert('''RGB'''							).resize((64, 64)							)
    if str(_lowercase							).startswith('''mps'''							):
     _UpperCamelCase:     Dict							= torch.manual_seed(_lowercase							)
    else:
     _UpperCamelCase:     Union[str, Any]							= torch.Generator(device=_lowercase							).manual_seed(_lowercase							)
    _UpperCamelCase:     List[str]							= {
        '''prompt''': '''A painting of a squirrel eating a burger''',
        '''image''': init_image,
        '''mask_image''': mask_image,
        '''generator''': generator,
        '''num_inference_steps''': 2,
        '''guidance_scale''': 6.0,
        '''output_type''': '''numpy''',
    }
    return inputs



   def   lowerCAmelCase						(	self					:  List[str]							):



    """simple docstring"""



    _UpperCamelCase:     Dict							= '''cpu'''  # ensure determinism for the device-dependent torch.Generator
    _UpperCamelCase:     str							= self.get_dummy_components()
    _UpperCamelCase:     Any							= StableDiffusionInpaintPipeline(**_lowercase							)
    _UpperCamelCase:     Any							= sd_pipe.to(_lowercase							)
    sd_pipe.set_progress_bar_config(disable=_lowercase							)

    _UpperCamelCase:     Tuple							= self.get_dummy_inputs(_lowercase							)
    _UpperCamelCase:     List[Any]							= sd_pipe(**_lowercase							).images
    _UpperCamelCase:     Optional[Any]							= image[0, -3:, -3:, -1]

    assert image.shape == (1, 64, 64, 3)
    _UpperCamelCase:     List[str]							= np.array([0.4727, 0.5735, 0.3941, 0.5446, 0.5926, 0.4394, 0.5062, 0.4654, 0.4476]							)

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



   def   lowerCAmelCase						(	self					:  Union[str, Any]							):



    """simple docstring"""



    super().test_inference_batch_single_identical(expected_max_diff=3E-3							)




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





   """simple docstring"""




   def   lowerCAmelCase						(	self					:  Dict							):



    """simple docstring"""



    super().tearDown()
    gc.collect()
    torch.cuda.empty_cache()



   def   lowerCAmelCase						(	self					:  Dict							):



    """simple docstring"""



    _UpperCamelCase:     Optional[Any]							= load_image(
        '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
        '''/sd2-inpaint/init_image.png'''							)
    _UpperCamelCase:     Tuple							= load_image(
        '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd2-inpaint/mask.png'''							)
    _UpperCamelCase:     Dict							= load_numpy(
        '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd2-inpaint'''
        '''/yellow_cat_sitting_on_a_park_bench.npy'''							)

    _UpperCamelCase:     Optional[Any]							= '''stabilityai/stable-diffusion-2-inpainting'''
    _UpperCamelCase:     Optional[Any]							= StableDiffusionInpaintPipeline.from_pretrained(_lowercase		,  safety_checker=_lowercase							)
    pipe.to(_lowercase							)
    pipe.set_progress_bar_config(disable=_lowercase							)
    pipe.enable_attention_slicing()

    _UpperCamelCase:     List[str]							= '''Face of a yellow cat, high resolution, sitting on a park bench'''

    _UpperCamelCase:     List[Any]							= torch.manual_seed(0							)
    _UpperCamelCase:     str							= pipe(
        prompt=_lowercase		,  image=_lowercase		,  mask_image=_lowercase		,  generator=_lowercase		,  output_type='''np'''		,  )
    _UpperCamelCase:     Optional[Any]							= output.images[0]

    assert image.shape == (512, 512, 3)
    assert np.abs(expected_image - image							).max() < 9E-3



   def   lowerCAmelCase						(	self					:  List[Any]							):



    """simple docstring"""



    _UpperCamelCase:     Optional[Any]							= load_image(
        '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
        '''/sd2-inpaint/init_image.png'''							)
    _UpperCamelCase:     Union[str, Any]							= load_image(
        '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd2-inpaint/mask.png'''							)
    _UpperCamelCase:     str							= load_numpy(
        '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd2-inpaint'''
        '''/yellow_cat_sitting_on_a_park_bench_fp16.npy'''							)

    _UpperCamelCase:     Optional[Any]							= '''stabilityai/stable-diffusion-2-inpainting'''
    _UpperCamelCase:     int							= StableDiffusionInpaintPipeline.from_pretrained(
        _lowercase		,  torch_dtype=torch.floataa		,  safety_checker=_lowercase		,  )
    pipe.to(_lowercase							)
    pipe.set_progress_bar_config(disable=_lowercase							)
    pipe.enable_attention_slicing()

    _UpperCamelCase:     Optional[Any]							= '''Face of a yellow cat, high resolution, sitting on a park bench'''

    _UpperCamelCase:     Dict							= torch.manual_seed(0							)
    _UpperCamelCase:     int							= pipe(
        prompt=_lowercase		,  image=_lowercase		,  mask_image=_lowercase		,  generator=_lowercase		,  output_type='''np'''		,  )
    _UpperCamelCase:     Optional[Any]							= output.images[0]

    assert image.shape == (512, 512, 3)
    assert np.abs(expected_image - image							).max() < 5E-1



   def   lowerCAmelCase						(	self					:  Any							):



    """simple docstring"""



    torch.cuda.empty_cache()
    torch.cuda.reset_max_memory_allocated()
    torch.cuda.reset_peak_memory_stats()

    _UpperCamelCase:     Union[str, Any]							= load_image(
        '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
        '''/sd2-inpaint/init_image.png'''							)
    _UpperCamelCase:     Optional[Any]							= load_image(
        '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd2-inpaint/mask.png'''							)

    _UpperCamelCase:     Optional[int]							= '''stabilityai/stable-diffusion-2-inpainting'''
    _UpperCamelCase:     Union[str, Any]							= PNDMScheduler.from_pretrained(_lowercase		,  subfolder='''scheduler'''							)
    _UpperCamelCase:     Optional[int]							= StableDiffusionInpaintPipeline.from_pretrained(
        _lowercase		,  safety_checker=_lowercase		,  scheduler=_lowercase		,  torch_dtype=torch.floataa		,  )
    pipe.to(_lowercase							)
    pipe.set_progress_bar_config(disable=_lowercase							)
    pipe.enable_attention_slicing(1							)
    pipe.enable_sequential_cpu_offload()

    _UpperCamelCase:     Optional[int]							= '''Face of a yellow cat, high resolution, sitting on a park bench'''

    _UpperCamelCase:     Union[str, Any]							= torch.manual_seed(0							)
    _UpperCamelCase:     Union[str, Any]							= pipe(
        prompt=_lowercase		,  image=_lowercase		,  mask_image=_lowercase		,  generator=_lowercase		,  num_inference_steps=2		,  output_type='''np'''		,  )

    _UpperCamelCase:     Dict							= torch.cuda.max_memory_allocated()
    # make sure that less than 2.65 GB is allocated
    assert mem_bytes < 2.65 * 10**9